Mixture of Parrots: Experts improve memorization more than reasoning
Authors: Samy Jelassi
Clara Mohri
David Brandfonbrener
Alex Gu
Nikhil Vyas
Nikhil Anand
David Alvarez-Melis
Yuanzhi Li
Sham M. Kakade
eran malach
Venue: ICLR 2025 Poster (2025)
Keywords: Mixture of Experts
memorization
reasoning
Primary area: foundation or frontier models, including LLMs

Abstract: The Mixture-of-Experts (MoE) architecture enables a significant increase in the total number of model parameters with minimal computational overhead. 
However, it is not clear what performance tradeoffs, if any, exist between MoEs and standard dense transformers.
In this paper, 
we show that as we increase the number of experts (while fixing the number of active parameters), the memorization performance consistently increases while the reasoning capabilities saturate. 


We begin by analyzing the theoretical limitations of MoEs at reasoning. We prove that there exist graph  problems that cannot be solved by any number of experts of a certain width; however, the same task can be easily solved by a dense model with a slightly larger width. 
On the other hand, we find that on memory-intensive tasks, MoEs can effectively leverage a small number of active parameters with a large number of experts to memorize the data. 
We empirically validate these findings on synthetic graph problems and memory-intensive closed book retrieval tasks. 
Lastly, we  pre-train a series of MoEs and dense transformers and evaluate them on commonly used benchmarks in math and natural language. 
We find that increasing the number of experts helps solve knowledge-intensive tasks, but fails to yield the same benefits for reasoning tasks.

Published as a conference paper at ICLR 2025
MIXTURE OF PARROTS MEMORIZATION MORE THAN REASONING
EXPERTS IMPROVE
Samy Jelassi*
Clara Mohri
David Brandfonbrener
Alex Gu MIT
Harvard University
Harvard University
Harvard University
Kempner Institute
Nikhil Vyas Harvard University
Nikhil Anand Harvard University
David Alvarez-Melis Harvard University
Yuanzhi Li Microsoft Research
Kempner Institute
Kempner Institute
Sham M. Kakade Harvard University
Eran Malach
Harvard University
Kempner Institute
Kempner Institute
ABSTRACT
The Sixture-of-Expertt (MoE) architecture enables a significant increase in the total number of model parameters with minimal computational overhead. However, it is not clear what performance tradeoffs, if any, exist between MoEs and standard dense transformers. In this paper, we show that we increase the number of experts (while fixing the number of active parameters), the memorization performance consistently increases while the reasoning capabilities saturate. We begin by analyzing the theoretical limitations of MoEs at reasoning. We prove that there exist graph problems that cannot be solved by any number of experts of certain width; however, the same task can be easily solved by a dense model with a slightly larger width. On the other hand, we find that on memory-intensive tasks, MoEs can effectively leverage a small number of active parameters with a large number of experts to memorize the data. We empirically validate these findings on synthetic graph problems and memory- inteesivee closed book retrieval tasks. Lastly, we pre-train series of MoEs and dense transformers and evaluate them on commonly used benchmarks in math and natural language. We find that increasing the number of experts helps solve knowledge-intensive tasks, but fails to yield the same benefits for reasoning tasks.
INTRODUCTION
The explosion in capabilities of large language models in recent years has largely been enabled by scaling their size, measured by the number of parameters in the model. In the standard Transformer architecture, scaling the number of parameters entails a proportional increase in computational cost, e.g. doubling the number of parameters requires doubling the number of floatingpooin operations (FLOPs), making training and inference more computational intensive. Mixture-of-Exxe Expeersss (MoE) were introduced
solution for this problem (Shazeer al., 2017: Lepikhin al., 2020; Fedus et al., 2022). MoEs replace the single MLP in each Transformer block with multiple MLPs (called experts), where each token is routed to a few experts based on a linear routing function. The number of parameters in the MoE layer therefore increases with the total number of experts, while the compute increases only with the number of "active" experts (i.c., the number of experts to which the token is routed to). This offers a promising option for scaling models: increase the number of experts instead of the model dimension or its depth. For this reason, MoEs have become very popular, and many frontier models today are based the MoE architecture (Achiam et al., 2023; Databricks, 2023; Anil et al., 2023; Dai et al., 2024; Jiang et al., 2024; Yang et al., 2024).
"Correspondence Samy Jelassi <ajel ass fas.. www.shaavaddeeee

---

Published as a conference paper at ICLR 2025
NLP (world knowledge)
NLP (commonsense)
MATH
53
44
F1 Acauray (86)
50
Acturacy (50
-22
Accuracy
46
14
14
8-
16M
40M
150M
560M
2.18
16M
40M
150M
509M
2.18
16M
159M
580M
2.18
Number of total parameters
Number of total parameters
Number of total parameters
Dense transformer
MoE (I8M active parameters)
MoBi (58M active parameters)
MoF. (200M active parameters)
(a) Evaluation: world knowledge
(b) Evaluation: commonsense
(c) Evaluation: math
world knowledge. We train series of dense transformers and MoEs on 65B tokens from a corpus essentially made of Fineweb-edu, Cosmopedia and Wikipedia (see Section 5 for details). We then evaluate the models several world knowledge benchmarks (e.g.. TriviaQA (Joshi et al., 2017), Natural Questions (Kwiatkowski et al., 2019)) and report the average Fl accuracy. Surprisingly, a fixed number of total parameters, MoES with substantially fewer active parameters approximately match the performance of dense models. This highlights the importance of experts in tasks that require memorization. (b) Evaluation: commonsense. Here we evaluate the aforementioned pre-trained models on natural language commonsense benchmarks (c.g., HellaSwag (Zellers al., 2019), WinoGrande (Sakaguchi al., 2021)). On these reasoning tasks, we observe that MoEs perform worse than dense models and more significant benefits are obtained by increasing the number of active parameters. (c) Evaluation: math. Here we train series of dense transformers and MoEs on 65B tokens from a corpus essentially made of Proof-Pile2 (Azerbayev al., 2023) (see Section for details). The results are consistent with the ones in (b): MoEs perform worse than dense models at equal number of total parameters.
Figure I: (a) Evaluation:
In this work we study whether MoE. indeed offers a "free-lunch", enabling gains in performance with no computational cost. Interestingly, we find that the benefit from MoEs greatly depends on the task at hand. We show that for reasoning-basee tasks, such as graph problems and mathematical reasoning, MoEs offer limited performance gains, and increasing the number of experts cannot compete with scaling the dimension (width) of the model. On the other hand, for memory-intensiv tasks, we show that scaling the number of experts is competitive with scaling standard "dense" MLPs.
To demonstrate these claims, we begin with a theoretical analysis of MoEs and dense models. We use communication-con compplllittt lower bounds to show that a single-layer MoE requires critical dimension to solve a simple graph connectivity problem, implying that MoEs offer no benefit for solving this problem and only consume unnecessary memory. On the other hand, we show that for pure memorization task, where the model only needs to "remember" an arbitrary set of examples, scaling the number of experts is equivalent to scaling the number of parameters in dense transformers, implying significant computational gain when fixing the number of active parameters (Section 3). We continue by experimentally validating these results, comparing MoEs against dense models on synthetic training data. We train these models on finding the shortest path in random graphs, where we show that MoE accuracy does not improve as we increase the number of experts, but that accuracy consistently increases with width for dense transformer (Figure 4b). Following this, we train different models on the task of memorizing a large phone-book. We demonstrate that MoEs excel in memorization, matching the performance of dense transformers with the same number of total parameters but with substantially less computational cost (Figure 4a).
Finally, we train dense transformers and MoEs on real datasets of mathematical reasoning and natural language, and perform intensive benchmarking of these models on a wide variety of downstream tasks. For memory-intene number of experts can match the performance of large dense models (Figure 1a). However, we show that for tasks that rely on reasoning, scaling the number of experts cannot compete with increasing the model dimension (Figures lb-lc). Moreover, Moßs exhibit some memorization behaviors when trained on math problems (Figure 5). Taken together, our results show that the gains from using MoEs depend greatly the nature of the training data and downstream task, and that while MoEs can improve performance in certain cases, sometimes increasing the effective size (width) of the model is unavoidable.
tasks, MoEs surprisingly have great advantage, where increasing the
2

---

Published as a conference paper at ICLR 2025
2 RELATED WORK
Mixture of Experts. Mixture-of-Expete (MoE) date back to the work of Jacobs et al. (1991) I); Jordan & Jacobs (1994). Shazeer et al. (2017); Fedus et al. (2022) were the first to scale this idea to deep learning and obtain ofte-of-thhe-rt models in machine translation. Since then, several works have improved their routing algorithms (Lepikhin et al., 2020; Lewis et al., 2021; Roller et al., 2021; Clark al., 2022; Zhou et al., 2022; Antoniak et al., 2023; Zhong et al., 2024), have improved their downstream performance after finetuning (Du et al., 2022: Zoph et al., 2022) or made their training and inference more efficient (Rajbhandari al., 2022: Gale al., 2023; Pan et al., 2024; Tan et al., 2024). However, only a few papers have studied the science of MoEs and their comparison with dense transformers. Clark et al. (2022); Krajewski et al. (2024) establish scaling laws for MoEs. Chen et al. (2022) design a specific classification problem where a model with multiple experts provably outperforms one with only one expert. Shazeer et al. (2017); Lepikhin et al. (2020); Artetxe et al. (2021); Lewis et al. (2021); Fedus et al. (2022); Du et al. (2022) show that given a fixed FLOP budget, MoEs are always better. However, these papers claim that on a per parameter basis, MoEs always seem comparatively worse than dense models. In this paper, temper this claim by showing that it depends on the nature of the task at hand: on reasoning tasks, we validate this claim but on memory-intensive tasks, equally-sizes MoEs perform well dense transformers.
Language models and memorization. Large language models (LLMs) store a considerable amount of knowledge in their parameters (Petroni ct al., 2019; Heinzerling & Imui. 2020). They memorize useful knowledge such as facts and commonsense (Zhao et al., 2024). Many works studied how memorization occurs in LLMs by developing tools to locate the knowledge in the model (Meng et al., 2022: Allen-Zhu & Li, 2023; Liu et al., 2024) or by tracking the training dynamics (Tirumala et al., 2022: Speicher et al., 2024). We draw inspiration from Allen-Zhu & Li (2023) and evaluate the memorization of our models by pre-training them on a mixture of datasets that includes Wikipedia, and at test time, evaluate them on world knowledge benchmarks, which essentially question answering tasks on Wikipedia facts. With respect to theoretical findings, Kim et al. (2023); Mahdavi et al. (2023); Madden et al. (2024); Nichani al. (2024) provide upper bounds the number of parameters needed for dense transformers to perform memorization tasks under various conditions.
Language models and reasoning. In recent years, transformer- bassd language models have displayed remarkable effectiveness in solving broad range of reasoning tasks. Specifically, the reasoning capabilities of transformers have been studied in the context of arithmetic problems (Jelassi et al., 2023; Cho et al., 2024; Hou et al., 2024; Zhou et al., 2024; McLeish et al., 2024; Lee el., 2023), mathematical reasoning (Zhang et al., 2022; Imani et al., 2023; Wei et al., 2022) graph problems (Sanford et al., 2024; Fatemi et al., 2023; Jin et al., 2023; Wang et al., 2024) and code challenges (Shi et al., 2024; Zhu et al., 2024). Recently, ofate-of-the-an art language models were used for solving complex math olympiad problems (DeepMind, 2024; NuminaMath, 2024; OpenAI, 2024). With respect theoretical findings, various works study the reasoning capabilities of transformers, relating their expressive power to other complexity classes and formal languages (Weiss et al., 2021; Zhou et al., 2023; Strobl et al., 2024; Chen & Zou, 2024). Other works study how chain-of-thoun can improve the reasoning capabilities of language models in terms of expressive power and learnability (Abbe al., 2024; Merrill & Sabharwal, 2023; Malach, 2024).
3 THEORY: REPRESENTATIONAL CAPACITY
In this section, we analyze the capability of MoE transformers compared to standard (dense) models. We begin by studying a simple graph problem that requires scaling the hidden dimension of the transformer, showing that MoEs with small hidden dimension cannot solve this problem, regardless of the number of experts used. Then, we show that MoËs can effectively memorize random inputs, requiring significantly less computational resources (active parameters) compared to dense models.
3.1 SETTING
Consider a lay-l transformer f E Transformer," H.I N and has logarithmic bit-precision. f embeds the input into dimension via the function Φ. has
m.. which takes as input a sequence of length
3

---

Published as a conference paper at ICLR 2025
H 2 attention heads, whose outputs combined via concatenation before apply point-wise function x
Define the parameters as Qh. Vh. Kn E Rmxm Φ & - R R R. The output of is:
f( σN) + (/softmax [[softmax(((xx) (d(NN) (aax)TTQuaa QuKK 6(X)) (X))V
( nE(H))
f is a dense transformer, if 1 is an MLP, i.e. function of the form:
(x)) = NouTool o( + b), for W R xm ; E IR'", E Rm
where a is the ReLU activation function. f E Transformer M,IKK,K is sparse (MoE) transformer with K experts if is function of the form:
(x)) = u, (W,w + b.) for argmaxr,
where W1,... W E Rm xm bi...., by RT ui..... RT are the parameters of each expert and r1 Th define the routing function (we use top-1 routing).
3.2 MoEs REQUIRE A CRITICAL HIDDEN SIZE TO SOLVE GR APH REASONING TASKS
In this section, we analyze the graph reasoning capabilities of dense and sparse transformers. We define the length-2 path problem graph, and use it means to understand other graph reasoning tasks such as graph connectivity, shortest path, and cycle detection.
Definition 3.1 (Length-2 2 Path Problem). The input is graph = (V; E). The source and destination d E V are fixed for all tasks as the and [V] vertex. The length-2 path problem asks whether there is path of length from 8 to d.
Graph connectivity, shortest path, and cycle detection are all graph reasoning tasks which reduce to the length-2 -2 path problem due to (Sanford et al., 2024) and Lemma D.2. We provide lower-bound on the width required for sparse transformer to solve the length-2 path problem, and an upper- on the width required for dense transformer to solve the problem. Further, between dense and sparse transformers with the same number of parameters: for a sufficiently large amount of experts in the sparse model, it cannot solve the same problem that dense model can solve with the same amount of total parameters.
show separation
Lower bound on width of depth-1 MoE for reasoning. We begin by showing a lower-bound on the width for a depth-1 mixture of expert model for the length-2 path problem. This lower bound implies a lower bound for search and retrieval tasks such as graph connectivity, shortest path, and eycle detection.
Theorem 3.2 (Length-2 path lower-bound sparse transformers). For some input sequence G (V, E), fxx two disjoint subsets A, B c [N - 1], and consider a single-layer transformer f E H.H.K.K with Olog N)-bit precision that solves length-2 path for any input X where A is a function of edges with the source s, XB a function of edges with the destination d. Then, f has width satisfying mHl = SQVVV log N).
Transformer N
The proof follows almost identically from the proof in (Sanford et al., 2024) for the class Transformer, M M1, The original proof does not place constraints on the function and is based on communication-compleitty argument. As such we may design vé that it first routes and then chooses which expert to apply. We give complete proof in Appendix D. As such, the result of (Sanford et al., 2024) can also be extended to the class Transformer m. H.I.K
Upper bound on width of depth-1 dense transformer for reasoning. In this section we give an upper bound for the width required for dense model solve the length-2 path problem.
In multi- Transformers, each layer outputs a vector of size m. However, since our focus in this section will be binary classification problems, we will let the transformer output single scalar, and we interpret the output of the final token as the prediction for the classification task.
4

---

Published as a conference paper at ICLR 2025
Theorem 3.3 (2ength-2 path width upper bound for transformer). There exists transformer of width [V], H = 1. and Ollog N) --it precision that solves length-2 path problem for any input.
The proof relies on an encoding of the inputs where the output values only exceed a certain threshold when u and the source and destination vertices, have edges with common vertex. We defer the proof to Appendix D.
Parameter-matched comparison dense and sparse depth-1 transformers. Using the lower- bound width required for sparse transformer (Theorem 3.2) and the upper-bound on width required for a dense transformer (Theorem 3.3), we compare dense and sparse transformers when they have the same number of total parameters. We find that when the number of experts exceeds (log N)?, the sparse model is unable to solve the same task as the dense model.
Corollary 3.4. Consider a sparse transformer (with experts) and a dense transformer with the same number of parameters. There exists number of experts K so that the the sparse model is not able to solve the reasoning task, but the dense transformer solves the task.
Proof. Suppose we have two depth-1 transformers, where one is a dense model and the other is mixture of experts with K experts. Let the width of the dense model be ma, and the width of the sparse model be m,. The number of parameters in the dense model is (m2) and the number of parameters in the sparse model is O(K In order to match the number of parameters, it must be the case that m,
ma Suppose we ma [V. as this sufficient solve the above problems.
VK"
For any K 2 22( R) the sparse model is not sufficiently wide solve the problem.
3.3 MoEs USE THEIR EXPERTS TO SOLVE MEMORY-INTENSIVE TASKS
In this section, we provide an upper-bound on the number of parameters necessary for sparse trans- former to solve memorization tasks, followed by a lower-bound on the number of parameters needed for a dense transformer to solve the same task, We use these results to compare the memorization capabilities of dense and sparse transformers with the same number of active parameters. We find that with enough experts, the sparse transformer is able to solve memorization tasks with less active parameters than the dense transformer. In both bounds we assume that transformer has logarithmic number bits to encode each parameter.
We consider sequences {(X', 3)){= where X E Nxm are input sequences of length N in dimension m such that []] is sampled from a Gaussian distribution (0 Lm)) We assume y1,...: J/N E {+1} arbitrary labels for the sequences. The objective is for a transformer to memorize these sequences, i.e. map each input X' to label y:. The classification is determined by the sign of the last token output.
Upper-bound MoE for memorization. We begin by showing that, with high probability the choice of the inputs, the MoE architecture can memorize (i.e., arbitrarily label the examples), with small number of active parameters.
Theorem 3.5. With probability at least 0.99, there exists a one-layer MoE transformer with K
experts, using Ô 60=+
K + Km)
+K)) active parameters and Ô (n + Km) total parameters stored in Õ())
bits that, when applied each sequence X3, oulputs at the last token value whose sign matches yo. i.e., sign(f(X,)) = Yi for all = I,...,
Specifically, if choose K = Vn/m m get that an MoE architecture solve the memorization problem with Ö(ynm) parameters. To prove this, we show that for a particular routing function, the number of samples routed to each expert is approximately n/ K. Then, we show that an expert with Ö(n/mK) neurons can memorize a sample of size ((n/ K). We present the proof in Appendix D.2.
Lower bound on memorization with dense Transformer. Next, we give lower-bound on the number of parameters for a dense transformer to perform memorization.
We use Ô and Ô hide logarithmic factors.
5

---

Published as a conference paper at ICLR 2025
Theorem 3.6 (Lower bound for dense model). Given the same task as above, dense Transformer requires ((n) parameters to solve the memorization task.
This bound follows from the fact that there are 2" possible labels for any fixed set n imputs, and at most 2oW functions with W parameters and bit per parameters. The proof is in Appendix D.2.
Separation between MoEs and Dense Models. Observe that the previous results on memorization imply a separation between MoEs and dense models in terms of the number of active parameters. Namely, we show that an MoE with (y/m)) active parameters can memorize, while a dense model requires (2)) parameters. So, for n m. MoEs are significantly more efficient. Comparing the number total parameters, MoEs require Ö(n + Km) parameters, so both MoE and dense models have linear dependence on in the total parameter count.
4 SYNTHETIC EXPERIMENTS
In the previous section, we proved that there exist graph connectivity problems that cannot be solved by any number of experts of certain width but the same task can be solved by a dense model with a slightly larger width. Our goal in this section is to verify that our theoretical analysis bears out experimentally when training models from scratch on synthetic data, before moving on to study pre-trained models in Section 5, We mainly focus on two tasks: the shortest path problem (Figure 2), which we use as a synthetic task to represent reasoning problems, and the phone-book task (Figure 3), to measure the recall ability of our models. Our experiments in this section highlight that adding experts yields greater performance improvements on memorization tasks than reasoning tasks.
Masked tolonas
Shoctest path
DDS
Outout
7
7
Z
Z
Model
Tapsst
DOS
Imput graph
Illustration of the shortest path task. We feed the model with sequence that lists all the edges in the input graph and ends with query (in green) which asks the model find a shortest path between two vertices (from vertex to vertex 4 in the figure). The model then autoregressively returns the shortest path (in purple).
Figure 2:
4.1 EXPERIMENTAL SETUP
Architecture. We opt for the Mistral (Jiang et al., 2023) and Mixtral (Jiang et al., 2024) architectures as the backbones of our Transformer and MoE models, respectively. The two architectures are identical and only differ by the addition of a gating module and multiple experts in Mixtral. For both model types, fix the number of layers to L 12. For the dense transformers,
vary model size by sweeping the width d E (256, 512, 1024]. For MoEs, we sweep over widths d' E (256, 512] and the number of experts E' {8, 16, 32, 64). To be consistent with experiments in Section 5, set the intermediate dimension in the FFN block to be equal to (and not 4d). We use token-choice routing, do not apply any token dropping and each token routed the top-2 experts. Lastly, in both this section and Section 5, we report for each model the number of non-embedding parameters which we refer to the total number of parameters.
Shortest path task. For a graph with n vertices, our token space V is of size + 6 with tokens encoding the vertices and some special tokens: V= = (1....: n. (ENGE) (mos), (Eos) (PAD). (sEP) ./) where (BOS) is the beginning of sentence token, (Fos) the end of sentence token, (PAD) the padding token, (EDGE) is the token indicating an edge between two vertices and, (SEr) and "/" are separator tokens. Each sequences describes the graph by a list of all the edges followed by two randomly sampled vertices and the shortest path between these latter (see Figure 2). All the graphs are directed and sampled according to the Pedös-Rénne model, with vertices and probability p for each edge to exist. We vary n E (25, 30, 50, 40, 45, 50, 55} and set such that the average length the shortest path is 3.5. Each train/test pair corresponds one value of (n. p), we do not mix graph sizes.

---

Published as a conference paper at ICLR 2025
Masloed tokens
Phone mumber
Name
KOS
abcde
60907642
Outguet
fghij 40652457
Model
klmno 20826343
pqrst 53143138
Iuput
8
uwwxy 80506194
BO8
SEP
Query
Figure 3: Illustration of the phone-book task for closed-book retrieval. The model is first trained to memorize a phone-book (illustrated on the right). Then, randomly select name in the phone-book (in green) and ask the model return their phone numher (in purple) without access the phone-hook.
PHONE-BOOK
SHORTEST PATH (n=50)
5.5M
80
Phonebook size
Accuracy (%)
75
1M
500k
68
100k
50k
63
5M
22M
88M
352M
700M
5M
22M
88M
352M
700M
Number of total parameters
Number of total parameters
Dense transformer
MoE (10M active parameters)
MoE (42M active parameters)
(a) Phone-book memorization
(b) Shortest path (50-node graphs)
Figure 4: (a) Phone-book memorization: train a series of dense transformers MoEs on phone-books of varying sizes and then evaluate their memorization capacity. We report the maximal phone-book size where the model obtains more than 90% accuracy. The maximal phone-book size correlates with the total (and not active) number of parameters. (b) Shortest path (total parameters): We train models to find the shortest path 50-node graphs and report the test accuracy. Here, increasing the number of experts provides limited improvements and the performance rather correlates with the number of active parameters.
Phone-book task. Our token space V is of size 39 and made of the alphabet letters, digits and special tokens: = {a..... ......,.. 9, (Bos) (uos). (sue)). We generate phone-books where the names consist of letters and the phone numbers of 8 digits (see Figure 3). We ensure that both the names and numbers are unique.
Datasets. For the graph experiments, the training set size is le6 and the test set consists of 1e3 held-out examples that sampled from the same distribution the training examples. For the phone-book experiments, we vary the training set size over {le5, 5e5, 1e6, 1.5e6, 2e6, 2.5e6, 3e6} and the test set consists of 1e3 queries from the training set.
Optimization. We use the AdamW optimizer (Loshchilov et al., 2017) with a weight decay equal to 0.1. We sweep the learning rate over (5e-5, 1e-4, 5e-4, le-3), the number of epochs over (2,5, 5. 10, 15), and set the maximal possible batch size among {8, 16, 32}. We use a warmup during the 20% first training steps and linear decay scheduler. All models prediction. In the graph task, we apply mask on the input instance so that we only penalize the model whenever it makes mistake the labels (and not on the inputs and labels jointly). In the phone-book experiment, we do not apply any masking.
trained by next-token
Evaluation. For each task we compute the exact accuracy, i.e. we count the generation as correct only it fully matches the ground truth. For the phone-book task, we report the size the maximal phone-book where we observe at least 90% exact accuracy.
4.2 MEMORIZATION TOTAL PARAMETERS PREDICT PERFORMANCE
We train dense transformers and MoEs on phone-books books of different sizes and at test time, evaluate whether they memorize the phone number of some names. Figure 4a reports the maximal phone-book size where the model manages to get an accuracy greater than 90%. This gives us an estimate of the memorization capacity of the model.
7

---

Published as a conference paper at ICLR 2025
The findings clear: no matter the number of active parameters, MoEs match the performance of dense transformers with the same number of total parameters. This suggests that MoEs are able to effectively leverage the extra parameters in additional experts by routing tokens to the experts that contain the necessary information from the training corpus. This scaling is remarkable in this case since it even holds when we only routing 2 out of 64 experts. For instance, we find that MoE model with only 42M active parameters outperforms a dense model with 10x as many parameters. This type of impressively efficient memorization capacity may be major reason behind the success of MoE architectures.
4.3 REASONING: TOTAL PARAMETERS DO NOT PREDICT PERFORMANCE
We train dense transformers and MoEs on the shortest path task and then query the models to find the shortest paths in novel, held-out graphs. Figure 4b reports the performance on graphs with 50 nodes with respect to their number of total parameters. Contrary to the phone-book experiment, increasing the number of experts does not consistently improve the performance of MoEs. Essentially, we find that active parameters rather than total parameters is better predictor of performance for these reasoning tasks.
To connect back to the theory from Section 3, note that active parameters is directly determined by the width of the network since we always route to exactly 2 experts and fix the depth. Thus, these results corroboaate the theory by showing that width (i.e. active parameters) determines the performance on these graph reasoning problems and that increasing the number of experts is not helpful. In Section 5, will further corroborate this idea through evaluation of mre-trained models commonsense and math reasoning benchmarks.
PRE-TRAINED MODELS
In this section, we pre-train dense transformers and MoEs and compare their performance on standard math and natural language benchmarks. We break the downstream tasks into those that require more memorization and those that require more reasoning. Here, memorization refers recall in that we measure the ability of the models to retrieve real-world facts, The memorization-intensse tasks test for "world knowledge" and consist of benchmarks like TriviaQA (Joshi et al., 2017). We break the -eatoning-intensive tasks into two subcategories; one for natural language reasoning tasks like WinoGrande (Sakaguchi et al., 2021) and another for mathematical reasoning tasks like Hendrycks-MAN MATH (Hendrycks et al., 2021). These tasks may be seen real-world analogs of the stylized phone- book and shortest path tasks studied in Section 4.
We observe that performance while performance on reasoning tasks depends more on the number of active parameters (Figure 1). Additionally, we conduct an experiment that indicates memorization from MoEs may be harming reasoning performance since there is a larger gap between train and test accuracy for MoEs than dense models fixed total parameters (Figure 5). Finally, we conduct an ablation where we compare models at fixed validation perplexity rather than model size. We find that MoEs perform better on world knowledge tasks and similarly on reasoning tasks compared to dense models (Figure 6).
world-knowledge tasks is governed by the total number of parameters
5.1 SETUP
Architecture. We train dense transformers and MoEs using the OLMoE codebase (Muen- nighoff et al., 2024). We set the number of layers L = 20 and vary the width d E (256, 512, 1024, 2048, 4096) for dense transformers and larly to Muennighoff et al. (2024), we consistently set the intermediate dimension in the FFN/MoE blocks to d (and not 4d). For MoEs, we vary the number of experts E E {8, 16, 32, 64). For the specific case of width 256, we also train MoE with 256 experts because its parameter count approx- imately matches the one of a width-2048 dense model and thus, we can compare the downstream performance of the two models. We use t2p-2 token-choice routing, without token dropping which is implemented in the dMoE function from the Megablocks package (Gale et al., 2023). We leave the study of MoEs trained with other routing mechanisms for future work.
(256, 512, 1024] for MoEs. Simi-
8

---

Published as a conference paper at ICLR 2025
Training hyperparameters. We use the AdamW optimizer (Loshchilov al., 2017) with weight decay equal to 0.1. We set the learning rate to 0.001, train on 63B tokens (60k steps) with batch size 512 and sequence length of 2048. We use warmup during the 20% first training steps and linear decay scheduler. We train our models using FSDP (Zhao et al., 2023).
Pre-training datasets. We train two collections of models, one collection on natural language and another one on math. The natural language dataset is mixture constituted of FineWeb-ee (Penedo et al., 2024), Cosmopedia (Ben Allal et al., 2024), Wikipedia and the training sets of the downstream tasks we evaluate on. The math dataset is a mixture made of Proof-Pile 2 (Azerbayev et al., 2023) and instruction datasets such as AapenMathInssen et al., 2023). Each of the two training mixture approximately totals 65B tokens. A precise description of the training mixtures can be found in Appendix B.
(Toshniwal et al., 2024) and MetaMathQA (Yu
Evaluation. We measure the validation perplexity 5,000 held-out sequences sampled from the training distribution. And we evaluate our models on a series of natural language and math benchmarks. Explicitly, we divide them into three categories:
World-knowledge tasks: TriviaQA (Joshi et al., 2017), Natural Questions (Kwiatkowski et al., 2019), HotpotQA (Yang et al., 2018), WebQuestions (Berant et al., 2013), WeeQusson (Talmor & Berant, 2018).
Commonsense tasks: ARC-C and ARC-E (Clark et al., 2018), Commonsense (Talmor et al., 2018), HellaSwag (Zellers et al., 2019), Openbook (Mihaylov et al., 2018), PIQA (Bisk et al., 2020), SciQ (Welbl et al., 2017), SIQA (Sap et al., 2019), WinoGrande (Sakaguchi et al., 2021). Math: SVAMP (Patel et al., 2021), 1). GSM8k (Cobbe et al., 2021), GSM-Hard (Gao et al., 2023), Mandrycks-MATHH (Hendrycks et al., 2021), Minerva-MATH (Lewkowycz et al., 2022).
In all our experiments, we plot the average accuracy for each of these three categories. We report the corresponding per-task performance in Appendix C.
GSM8k
30
Gener Ceneralization gap
25
18
14
5.2 RESULTS
7
Experts improve memorization more than reasoning.
3
We observe that the conclusions from our theoretical results and synthetic experiments also hold when pre-training and evaluating language models natural language and math. In Figure la, we report the accuracy of our models with re- spect to the number of total parameters. All the lines in the plot approximately coincide which implies that regardless of the number of active parameters, MoEs can effectively use their routing to leverage all of their parameters to solve memory-intensive tasks. On the other hand, on common-
16M
49M
159M
569M
2.1B
Number total parameters
(a)
HENDRYCKS MATH
Generalization gap
37
sense and math benchmarks (Figures Ib, we find that MoEs do not reach the performance of dense models with
8
2-
16M
49M
159M
569M
2.1B
the same number of total parameters. This indicates that for these reasoning tasks, increasing the dense model width is more effective that adding experts.
Number of total parameters
(b)
Dense transformer
MoE (18M active parameters)
On math tasks, MoEs display a higher train-test gap than dense models, suggestive of memorization. We provide additional evidence that memorization occurs in pre-trained MoEs by considering the generalization gap. Figure 5 we select 6.319 random problems from the OphhnMtthlnstrutt dataset, which is part of the training mixture data. More precisely, we pick 5,000 Hendrycks-MATHH like examples and 1,319 GSM8k-like examples to ensure that the number of training examples matches with the corresponding number of examples in GSM8k and Hendrycks- -MATH test sets. We then report the generalization gap, which is the gap between the accuracy on training examples and test examples. While both dense transformers and MoEs make a single pass on the OpenMatthnntrrut dataset, Figure 5 shows that at scales beyond 1.5M parameters, MoEs suffer from more significant generalization gap than
MoE (58M active parameters)
MoE (200M active parameters)
Figure 5: Generalization gap when the test ist GSM8k and Hendrycks-MATH (b).

---

Published as a conference paper at ICLR 2025
NIP (world knowledge)
(commonsense)
MATH
53
44-
F1 Acauray (8)
20-
Accuracy (%)
50
Accuracy (1Q)
32
46
14-
42
39
6.1
7.3
12.2
0.2
12.2
3.2
3,9
Validation perplexity
Validation perplexity
Validation perplexity
Dense transformer
MoE (I8M active parameters)
MoE (58M active parameters)
Moli (200M active parameters)
(a)
(b)
(c)
Figure 6: (a) On world knowledge benchmarks, MoEs consistently outperform dense transformers in
downstream performance when fixing the validation perplexity. (b-c) In reasoning benchmarks, dense transformers perform about the same MoEs fixed validation perplexity. MoEs can achieve these perplexities with less active parameters, but may require substantially more total parameters.
dense transformers. This is suggestive that MoEs are more prone to overfit to the problems they have been pre-trained on than dense models.
MoE models excel at world knowledge tasks but match dense models in reasoning when perplex- ity is fixed. Finally, we focus on the relationship between validation perplexity and downstream performance in Figure 6. Rather than comparing models by their parameter count, we can compare them based on how well they fit the training distribution as measured by validation perplexity. Even though two models may have the same perplexity. they will have learned different functions. The question is then if we can see any high level patterns in which types of functions particular model class more likely to learn. Figure 6a shows that at a fixed perplexity, the MoE models outperform the dense models on world knowledge tasks.
On the other hand, Figures 6b and 6c show that MoEs and dense models perform about the same on the reasoning tasks at fixed validation perplexity. Note that both dense Transformers and Moïs are trained with the objective of minimizing the perplexity (or loss). However, there could be multiple strategies to achieve the same loss, e.g. by memorizing pieces of data or by improving "reasoning" capabilities. Which strategy the model prefers is determined by the implicit bias of the architecture. factual recall, MoEs achieve better accuracy for
In the experiments, we see that for memorization/ / the same pereplexity value, suggesting that they "prioritize" memorization over other capabilties.
6 DISCUSSION
In recent years, scaling up the number of parameters in Transformers has been the dominant approach for improving performance on language modeling. A standard Transformer of dimension d and sequence length T, has number of parameters which scales with Odd), and run-time that scales with L²). Improving the efficiency can either attempt to reduce the dependence on L or d. Sub-quadratic attention variants attempt to improve dependence on L (Katharopoulos al., 2020; Peng et al., 2023; Fu et al., 2022; Gu & Dao, 2023), while MoEs attempt to improve dependence on by scaling the number of parameters without scaling the dimension of the model.
This paper illuminates the costs and benefits of this reduced dependence on d. We show that for some reasoning-intern intenssive tasks increasing the dimension d is inevitable, and scaling the computation with (d2)) seems unavoidable. This remains true regardless of the different design choices in the MoE architecture and is backed up empirically. There is increasing interest in developing non-MoE models with sub-quadratic dependence d, using some structural assumptions on the weight layers (Kamalakara et al., 2022; Dao et al., 2021; 2022; Fu et al., 2024), which could provide an alternative.
On the other hand, we find that MoEs are highly effective at knowledge intensive tasks. They are able to much more efficiently memorize facts than dense models with a similar number of active parameters, even matching the performance of dense models with the same number of total parameters. This suggests that MoEs are valuable as memorization machines and perhaps this particular capability can be leveraged while relying on other architectures for more ceasoning-intense tasks.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENTS
We thank Cyril Zhang for helpful discussions and Max Shad for his support while running the experiments. Kempner Institute computing resources enabled this work. Samy Jelassi acknowledges funding supported by the Center of Mathematical Sciences and Applications. Alex Gu is supported by the National Science Foundation (NSF) Graduate Research Fellowship under Grant No. 2141064. David Alvarez-Melis acknowledges support from the Kempner Institute, the Aramont Fellowship Fund, and the FAS Dean's Competitive Fund for Promising Scholarship. This work has been made possible in part by a gift from the Chan Zuckerberg Initiative Foundation to establish the Kempner Institute for the Study of Natural and Artificial Intelligence; support from the Office of Naval Research under award N00014-22-1-2371 and the National Science Foundation Grant under award #IIS 2229881.
REFERENCES
Emmanuel Abbe, Samy Bengio, Aryo Lotfi, Colin Sandon, and Omid Saremi. How far can transformers reason? the locality barrier and inductive scratchpad, 2024. URL https:
warxiv.orgg/oooooooooooo. abs 2406. 06467.
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Iige Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.00 08774, 2023.
Zeyuan Allen-Zhu and Yuanzhi Li. Physics of language models: Part 3.1, knowledge storage and extraction. arXiv preprini arXiv:2309. 14316, 2023.
Rohan Anil, Sebastian Borgeaud, Yonghui Wu, Jean-Baptiste Alayrac, Jiahui Yu, Radu Soricut, Johan Schalkwyk, Andrew M Dai, Anja Hauth, Katie Millican, et al. Gemini: A family of highly capable multimodal models. arXiv preprint arXiv:2312. 1805, 1, 2023.
Szymon Antoniak, Sebastian Jaszezur, Michal Krutul, Maciej Pióro, Jakub Krajewski, Jan Ludziejew- ski, Tomasz Odrzygóždú, and Marek Cygan. Mixture of tokens: Efficient 11ms through cross- example aggregation. arXiv preprint arXiv: 10. 15961, 2023.
Mikel Artetxe, Shruti Bhosale, Naman Goyal, Todor Mihaylov, Myle Ott, Sam Shleifer, Xi Victoria Lin, Jingfei Du, Srinivasan Iyer, Ramakanth Pasunuru, al. Efficient large scale language modeling with mixtures of experts. arXiv preprint arXiv:2112.10663 2021.
Zhangir Azerbayev, Hailey Schoelkopf, Keiran Paster, Marco Dos Santos, Stephen McAleer, Albert Q. Jiang, Jia Deng, Stella Biderman, and Sean Welleck. Llemma: An open language model for mathematics, 2023.
Loubna Ben Allal, Anton Lozhkov, Guilherme Penedo, Thomas Wolf, and Leandro von Werra. Cosmopedia, 2024. URL htttp:://huginggee cosmopedia.
hoodaaaseess /Hugg IgFacceB//
Jonathan Berant, Andrew Chou, Roy Frostig, and Perey Liang. Semantic parsing on freebase from question-answer pairs. In Proceedings of the 2013 conference empirical methods in natural language processing, pp. 1533-1544, 2013.
Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, al. Piqa: Reasoning about physical commonsense in natural language. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pp. 7432-7439, 2020.
Xingwu Chen and Difan Zou. What can transformer learn with varying depth? case studies on sequence learning tasks. arXiv preprint arXiv:2404.01601, 2024.
Zixiang Chen, Yihe Deng, Yue Wu, Quanquan Gu, and Yuanzhi Li. Towards understanding mixture of experts in deep learning. arXiv preprint arXiv:2208. 02813, 2022.
Hanseul Cho, Jaeyoung Cha, Pranjal Awasthi, Srinadh Bhojanapalli, Anupam Gupta, and Chul- hee Yun. Position coupling: Leveraging task structure for improved length generalization of transformers. arXiv preprini Stt::2005.206/1 2024.

---

Published as a conference paper at ICLR 2025
Aidan Clark, Diego de Las Casas, Aurelia Guy, Arthur Mensch, Michela Paganini, Jordan Hoffmann, Bogdan Damoc, Blake Hechtman, Trevor Cai, Sebastian Borgeaud, et al. Unified scaling laws for routed language models. In International conference machine learning, pp. 4057-0044400 4086. PMLR, 2022.
Peter Clark. Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try are, the ai2 reasoning challenge. arXiv preprint arXiv: /803.085300 2018.
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 preprint arXiv:21 10. 14168, 2021.
Damai Dai, Chengqi Deng, Chenggang Zhao, RX Xu, Huazuo Gao, Deli Chen, Jiashi Li, Wangding Zeng, Xingkai Yu, Y Wu, et al. Deepseekmoe: Towards ultimate expert specialization in mixture- of-experts language models. arXiv preprint arXiv:2401.0641 2024.
Amit Daniely, Memorizing gaussians with no weer-parameterizaion via gradient decent on neural networks. arXiv preprini arXiv:2003./2895 2020.
Tri Dao, Beidi Chen, Kaizhao Liang, Jiaming Yang. Zhao Song, Atri Rudra, and Christopher Re. Pixelated butterfly: Simple and efficient sparse training for neural network models. arXiv preprint arXiv:2112.0002 2021.
Tri Dao, Beidi Chen, Nimit S Sohoni, Arjun Desai, Michael Poli, Jessica Grogan, Alexander Liu, Aniruddh Rao, Atri Rudra, and Christopher Ré. Monarch: Expressive structured matrices for efficient and accurate training. In International Conference on Machine Learning, pp. 4690-4721. PMLR, 2022.
Databricks. Databricks Blog, 2023.
Introducing
dbrx:
A new state-of-the-art open 11m. URL /ttps::/www databricke com/blog/
Accessed: 2023-10-12.
DeepMind. Ai achieves silver-medal standard solving international mathemati-
cal olympiad problems. ai ai-solves-imo- www.eee--ooopooooeee-----eeeeee-
https://deepmind. google wtw.d/ooogollllleeeeeeeeeee /ttpiiiiccvvoboooo/ce/e/o
2024.
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 mixture-of-experts. In International Conference Machine Learning, pp. 5547-5569 PMLR, 2022.
Bahare Fatemi, Jonathan Halcrow, and Bryan Perozzi. Talk like a graph: Encoding graphs for large language models. arXiv preprint arX1v:2310.045500, 2023.
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(120):1-39, 2022.
Dan Fu, Simran Arora, Jessica Grogan, Isys Johnson, Evan Sabri Eyuboglu, Armin Thomas, Benjamin Spector, Michael Poli, Atri Rudra, and Christopher Ré. Monarch mixer: A simple sub-quadratic gemm-based architecture. Advances in Neural Information Processing Systems, 36, 2024.
Daniel Y Fu, Tri Dao, Khaled K Saab, Armin W Thomas, Atri Rudra, and Christopher Ré. Hungry hungry hippos: Towards language modeling with state space models. arXiv preprint arXiv:2212. 14052, 2022.
Trevor Gale, Deepak Narayanan, Cliff Young, and Matei Zaharia. Megablocks: Efficient sparse training with mfxture-of-experts, Proceedings of Machine Learning and Systems, 1:288-304, 2023.
Luyu Gao, Aman Madaan, Shuyan Zhou, Uri Alon, Pengfei Liu, Yiming Yang, Jamie Callan, and Graham Neubig. Pal: Program-aide -idedd language models. In International Conference on Machine Learning, pp. 10764-10799. PMLR, 2023.
12

---

Published as a conference paper at ICLR 2025
Albert Gu and Tri Dao. Mamba: Ainear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752, 2023.
Benjamin Heinzerling and Kentaro Inui. Language models as knowledge bases: On entity representa- tions, storage capacity, and paraphrased queries. arXiv preprint urXiv:2008.09005 2020.
Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. arXiv preprint arXiv:2103.0333 2021.
Kaiying Hou, David Brandfonbrener, Sham Kakade, Samy Jelassi, and Eran Malach. Universal length generalization with turing programs. arXiv preprint arXiv:2407.03310, 2024.
Shima Imani, Liang Du, and Harsh Shrivastava. Mathprompter: Mathematical reasoning using large language models. arXiv preprint arXiv:2303.05398 2023.
Robert A Jacobs, Michael I Jordan, Steven J Nowlan, and Geoffrey E Hinton. Adaptive mixtures of local experts. Neural computation, 3(1):79-87, 1991.
Samy Jelassi, Stéphane d' AAscoli, Carles Domingo-Enricn Yuhuai Wu, Yuanzhi Li, and François Charton. Length generalization in arithmetic transformers. arXiv preprint arXiv:2306. 15400, 2023.
Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, et al. Mistral 7b. arXiv preprint arXiv:2310.0682) 2023.
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 preprint S.Xiv22401.040000 2024.
Bowen Jin, Gang Liu, Chi Han, Meng Jiang, Heng Ji, and Jiawei Han. Large language models on graphs: A comprehensive survey, arXiv preprint arXiv:2312.027833 2023.
Michael I Jordan and Robert A Jacobs. Hierarchical mixtures experts and the em algorithm. Neural computation, 6(2):181-214, 1994.
Mandar Joshi, Eunsol Choi, Daniel S Weld, and Luke Zettlemoyer. Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension. arXiv preprint arXiv: orXiv:/705.000 1705.03551, 2017.
Siddhartha Rao Kamalakara, Acyr Locatelli, Bharat Venkitesh, Jimmy Ba, Yarin Gal, and Aidan N Gomez. Exploring low rank training of deep neural networks. arXiv preprint arXiv:2209. 2022.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mnns: Fast autoregressive transformers with linear attention. In Internationna conference machine learning, pp. 5156-5165. PMLR, 2020.
Junghwan Kim, Michelle Kim, and Barzan Mozafari. Provable memorization capacity of transformers. In The Eleventh International Conference on Learning Representations, 2023.
Jakub Krajewski, Jan Ludziejewski, Kamil Adamczewski, Maciej Pióro, Michal Krutul, Szymon Antoniak, Kamil Ciebiera, Krystian Król, Tomasz Odrzygóždž, Piotr Sankowski, et al. Scaling laws for fine-graine grained mixture of experts. arXiv preprint arXiv:2402.07871, 2024.
Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Jacob Devlin, Kenton Lee, et al. Natural questions: a benchmark for question answering research. Transactions of the Association for Computational Linguistics, 7:453-466, 2019.
B. Laurent and P. Massart. Adaptive estimation of quadratic functional by model selection. The Annals of Statistics, 132 2000. ISSN 00905364, 21688966. URL http: // www. . jator wtt.:ooo/r//saabe/ee///////
13

---

Published as a conference paper at ICLR 2025
Nayoung Lee, Kartik Sreenivasan, Jason D Lee, Kangwook Lee, and Dimitris Papailiopoulos. Teaching arithmetic to small transformers. arXiv preprint arXiv:2307.03381, 2023.
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. Gshard: Scaling giant models with conditional computation and automatic sharding. arXiv preprini arXiv:2006. 1122000666660 16668, 2020.
Mike Lewis, Shruti Bhosale, Tim Dettmers, Naman Goyal, and Luke Zettlemoyer. Base layers: Simplifying training of large, sparse models. In International Conference on Machine Learning, pp. 6265-6274. PMLR, 2021.
Aitor Lewkowyez, Anders Andreassen, David Dohan, Ethan Dyer, Henryk k Michalewski, Vinay Ra- masesh, Ambrose Slone, Cem Anil, Imanol Schlag, Theo Gutman-Solo, et al. Solving quantitative reasoning problems with language models. Advances in Neural Information Processing Systems, 35:3843-3857, 2022.
Bingbin Liu, Sebastien Bubeck, Ronen Eldan, Janardhan Kulkarni, Yuanzhi Li, Anh Nguyen, Rachel Ward, and Yi Zhang. Tinygsm: achieving; 80% on gsm8k with small language models. arXiv preprint arXiv:2312.09241,, 2023.
Yan Liu, Yu Liu, Xiaokang Chen, Pin- Yu Chen, Daoguang Zan, Min- Yen Kan, and Tsung- Yi Ho. The devil is in the neurons: Interpreting and mitigating social biases in language models. In The Twelfth International Conference on Learning Representations, 2024.
Iyya Loshchilov, Frank Hutter, et al. Fixing weight decay regularization in adam. arXiv preprint 17X1v:1711.05101, 5, 2017.
Liam Madden, Curtis Fox, and Christos Thrampoulidis. Upper and lower memory capacity bounds of transformers for next-token prediction. arXiv preprint arXiv:2405.137/1 2024.
Sadegh Mahdavi, Renjie Liao, and Christos Thrampoulidis. Memorization capacity of multi-head attention in transformers. arXiv preprint arXiv:2306.000 02010, 2023.
Eran Malach. Auto-regressive next-token predictors are universal learners. In Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 34417-34431. PMLR, 21-27 Jul 2024.
Sean McLeish, Arpit Bansal, Alex Stein, Neel Jain, John Kirchenbauer, Brian R. Bartoldson, Bhavya Kailkhura, Abhinav Bhatele, Jonas Geiping, Avi Schwarzschild, and Tom Goldstein. Transformers can do arithmetic with the right embeddings, 2024. URL https: /ttp///wwww..////////00/000 17399.
Kevin Meng, David Bau, Alex Andonian, and Yonatan Belinkov. Locating and editing factual associations in gpt. Advances in Neural Information Processing Systems, 35:17359-17372, 2022.
William Merrill and Ashish Sabharwal. The expresssive power of transformers with chain of thought. arXiv preprint arXiv:2310.0792. 2023.
Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can suit of armor conduct electricity? a new dataset for open book question answering. arXiv preprint arXiv: arXiv:/809.000 1809.02789, 2018.
Arindam Mitra, Hamed Khanpour, Corby Rosset, and Ahmed Awadallah. Orea-math: Unlocking the potential of slms in grade school math, 2024.
Niklas Muennighoff, Luca Soldaini, Dirk Groeneveld, Kyle Lo, Jacob Morrison, Sewon Min, Weijia Shi, Pete Walsh, Oyvind Tafjord, Nathan Lambert, Yuling Gu, Shane Arora, Akshita Bhagia, Dustin Schwenk, David Wadden, Alexander Wettig. Binyuan Hui, Tim Dettmers, Douwe Kiela, Ali Farhadi, Noah A. Smith, Pang Wei Koh, Amanpreet Singh, and Hannanch Hajishirzi. Olmoe: Open mixtureoofeexperts language models, 2024. URL https: : / /arxivi arxiv.orrs abs / 2409 02060.
Eshaan Nichani, Jason D Lee, and Alberto Bietti. Understanding factual recall in transformers via associative memories. arXiv preprint 24X 2024.
14

---

Published as a conference paper at ICLR 2025
NuminaMath. How numinamath won the 1st aimo progress prize. https://huggiiggfaeee wwwo-poorrss---p.-....... 2024.
OpenAl. Introducing openai ol. https: /openai com/ o2/, 2024.
Bowen Pan, Yikang Shen, Haokun Liu, Mayank Mishra, Gaoyuan Zhang, Aude Oliva, Colin Raffel, and Rameswar Panda. Dense training, sparse inference: Rethinking training of mixture-of-expert language models. arXiv preprint urXiv:2404,05567, 2024.
Keiran Paster, Marco Dos Santos, Zhangir Azerbayev, and Jimmy Ba. Openwebmath: An open dataset of high-quality mathematical web text. arXiv preprint arXiv:23 2023.
Arkil Patel, Satwik Bhattamishra, and Navin Goyal. Are nip models really able to solve simple math word problems? arXiv preprint arXiv:2 2021.
Guilherme Penedo, Hynek Kydlíček, Anton Lozhkov, Margaret Mitchell, Colin Raffel, Leandro Von Werra, Thomas Wolf, et al. The fineweb datasets: Decanting the web for the finest text data at scale. arXiv preprint a7X55 2024.
Bo Peng. Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, et al. Rwkv: Reinventing mns for the transformer era. arXiv preprint arXiv:2305. 13048, 2023.
Fabio Petroni, Tim Rocktäschel, Patrick Lewis, Anton Bakhtin, Yuxiang Wu, Alexander H Miller, and Sebastian Riedel. Language models as knowledge bases? arXiv preprint 19X9v: 2019.
Samyam Rajbhandari, Conglong Li, Zhewei Yao. Minjia Zhang, Reza Yazdani Aminabadi, Am- mar Ahmad Awan, Jeff Rasley, and Yuxiong He. Deepspeed-moe: Advancing mixture-of-expere inference and training to power next-generation ai scale. In International conference on machine learning, pp. 18332-18346. PMLR, 2022.
Stephen Roller, Sainbayar Sukhbaatar, Jason Weston, al. Hash layers for large sparse models. Advances in Neural Information Processing Systems, 34:17555-17566, 2021.
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale. Communications of the ACM, 64(9):99-106,
2021.
Clayton Sanford, Bahare Fatemi, Ethan Hall, Anton Tsitsulin, Mehran Kazemi, Jonathan Halcrow, Bryan Perozzi, and Vahab Mirrokni. Understanding transformer reasoning capabilities via graph algorithms. arXiv preprint arXiv:2405. 18512, 2024.
Maarten Sap, Hannah Rashkin, Derek Chen, Ronan LeBras, and Yejin Choi. Socialiqa: Commonsense reasoning about social interactions. arXiv preprint arXiv: 09728, 2019.
David Saxton, Edward Grefenstette, Felix Hill, and Pushmeet Kohli. Analysing mathematical reasoning abilities of neural models. arXiv preprint arXiv: 0904.9904...3 2019.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-exper layer. arXiv preprint
06538, 2017.
Quan Shi, Michael Tang, Karthik Narasimhan, and Shunyu Yao. Can language models solve olympiad programming? arXiv preprint arXiv:2404. 10952, 2024.
Till Speicher, Aflah Mohammad Khan, Qinyuan Wu, Vedant Nanda, Soumi Das, Bishwamittra Ghosh, Krishna P Gummadi, and Evimaria Terzi. Understanding the mechanics and dynamics of memorisation in large language models: A case study with random strings. 2024.
Lena Strobl, William Merrill, Gail Weiss, David Chiang, and Dana Angluin. What formal languages can transformers express? a survey. Transactions of the Association for Computational Linguistics, 12:543-561, 2024.
15

---

Published as a conference paper at ICLR 2025
Alon Talmor and Jonathan Berant. The web knowledge- base for answering complex questions. arXiv preprint arXiv: 1rX3::8803006433 2018.
Alon Talmor, Jonathan Herzig, Nicholas Lourie, and Jonathan Berant. Commonsensega: A question answering challenge targeting commonsense knowledge. arXiv preprint arXiv:1811.009) 2018.
Shawn Tan, Yikang Shen, Rameswar Panda, and Aaron Courville. Scattered mixture-of-expprts implementation. arXiv preprint arXiv:2403.082-5. 2024.
Kushal Tirumala, Aram Markosyan, Luke Zettlemoyer, and Armen Aghajanyan. Memorization without overfitting: Analyzing the training dynamics of large language models. Advances in Neural Information Processing Systems, 35:38274-38290, 2022.
Shubham Toshniwal, Ivan Moshkov, Sean Narenthiran. Daria Gitman, Fei Jia, and Igor Git-
man. Openmathinstruct-I: A 1.8 million math instruction tuning dataset. arXiv preprint arXiv:2402.10176, 2024.
Heng Wang. Shangbin Feng, Tianxing He, Zhaoxuan Tan, Xiaochuang Han, and Yulia Tsvetkov. Can language models solve graph problems in natural language? Advances in Neural Information Processing Systems, 36, 2024.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny Zhou, et al. Chain-of-thouuhh prompting elicits reasoning in large language models. Advances in neural information processing systems, 35:24824-24837, 2022.
Gail Weiss, Yoav Goldberg, and Eran Yahav. Thinking like transformers. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Re- search, pp. 11080-11090. PMLR, 2021. URL http: http::/procceeiing wttpssml...prrr wttp./gs.......e./ssss/sggee /v39/ weiss2la. html.
Johannes Welbi, Nelson F Liu, and Matt Gardner. Crowdsourcing multiple choice science questions. arXiv preprint arXiv: / 170700620, 2017.
An Yang, Baosong Yang, Binyuan Hui, Bo Zheng. Bowen Yu, Chang Zhou, Chengpeng Li, Chengyuan Li, Dayiheng Liu, Fei Huang, et Qwen2 technical report. arXiv preprint arXiv:2407.10671, 2024.
Zhilin Yang, Peng Qi, Saizheng Zhang, Yoshua Bengio, William W Cohen, Ruslan Salakhutdinov, and Christopher D Manning. Hotpotqa: A dataset for diverse, explainable multi-hop question 09600, 2018
answering, arXiv preprint
Longhui Yu, Weisen Jiang, Han Shi, Jincheng Yu, Zhengying Liu, Yu Zhang, James Kwok, Zhenguo Li, Adrian Weller, and Weiyang Liu. Metamath: Bootstrap your own mathematical questions for large language models. arXiv preprint 2rXiv:2309.12284, 2023.
Xiang Yue, Tuney Zheng, Ge Zhang, and Wenhu Chen. Mammoth2: Scaling instructions from the web. arXiv preprint arXiv:2405.03544 2024.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a machine really finish your sentence? arXiv preprint arXiv:/905.070 2019.
Yi Zhang, Arturs Backurs, Sébastien Bubeck, Ronen Eldan, Suriya Gunasekar, and Tal Wagner. Unveiling transformers with lego: synthetic reasoning task. arXiv preprint 2rXiv:2206.04301, 2022.
Yifan Zhang. Stackmathqa: A curated collection of million mathematical questions and an- swers sourced from stack exchange, 2024. URL https: httpss:///hgggggaceeee.. co/ co/datsetss math-ai/St wwt.aa/ttaaaaaaaaaaaaaaaaaa
Yanli Zhao, Andrew Gu, Rohan Varma, Liang Luo, Chien-Chin Huang, Min Xu, Less Wright, Hamid Shojanazeri, Myle Ott, Sam Shleifer, et al. Pytorch fsdp: experiences on scaling fully sharded data parallel. arXiv preprint ar0iv:2304. /277, 2023.
16

---

Published as a conference paper at ICLR 2025
Zirui Zhao, Wee Sun Lee, and David Hsu. Large language models commonsense knowledge for large-scale task planning. Advances in Neural Information Processing Systems, 36, 2024.
Zexuan Zhong, Mengzhou Xia, Danqi Chen, and Mike Lewis. Lory: Fully differentiable mixture- of-experts for autoregressive language model pre-training. arXiv preprint MxXiv:2005.03333 2024.
Hattie Zhou, Arwen Bradley, Etai Littwin, Noam Razin, Omid Saremi, Josh Susskind, Samy Bengio, and Preetum Nakkiran. What algorithms can transformers learn? study in length generalization, 2023. URL https:
org/ aatgg//b//aa /aa/aabs////0/ 16028.
Yanqi Zhou, Tao Lei, Hanxiao Liu, Nan Du, Yanping Huang, Vincent Zhao, Andrew M Dai, Quoc V Le, James Laudon, et al. Mixture-of-experts with expert choice routing. Advances in Neural Information Processing Systems,
2022.
Yongchao Zhou, Uri Alon, Xinyun Chen, Xuezhi Wang, Rishabh Agarwal, and Denny Zhou. Trans- formers can achieve length generalization but not robustly. arXiv preprint arXiv:2402.09371, 2024.
Qihao Zhu, Daya Guo, Zhihong Shao, Dejian Yang, Peiyi Wang, Runxin Xu, Wu, Yukun Li, Huazuo Gao, Shirong Ma, et al. Deepscck-coderr-2. Breaking the barrier of closed-source models in code intelligence. arXiv preprint arXiv:2406. 11931, 2024.
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 preprint arXiv:220..08906, 2022.
17

---

Published as a conference paper at ICLR 2025
A LIMITATIONS AND FUTURE WORK
While we provide substantial experiments on wide range of tasks, many these tasks are not pure memorization nor pure reasoning. We study these two extreme cases to better convey our message but it would be interesting in future work to understand the difference between dense and sparse models on tasks that mix both memorization and recall. Our pre-trained models have up S 2.1B parameters, but we recognize that large scale MoEs like Mixtral (Jiang et al., 2024), DeepSeek-V V2 (Dai et al., 2024), and others have orders of magnitude more parameters. We hypothesize that our results would still be meaningful at larger scales due to the strong theoretical underpinning, but it is not guaranteed. Moreover, as suggested above, would an interesting direction for future work to propose new architectures with reduced d dependence that can get the best of both worlds and solve reasoning and memorization tasks.
B DETAILS ON THE
DATASETS
In Section 5, we pretrain two collections of models, one on "natural language" and the other on "math". Here, we give precise breakdown of our training mixtures. We start with the "natural language" training mixture that totals 64B tokens:
37B tokens from Fineweb-edu dedup (Penedo al., 2024).
14B tokens from Cosmopedia (Ben Allal al., 2024).
12B tokens from Wikipedia (we loop over Wikipedia 3 times).
1B tokens from the training set of the downstream tasks we test on. We create 3 copies of each of these increase their presence in the mixture. The presence of these datasets is pretty important as argued in Allen-Zhu & Li (2023) so that the model is familiar with the downstream tasks at test time.
ComplexWebuuestinn WebQuestions training set (Talmor & Berant, 2018)
HotPorQA training set (Yang et al., 2018)
Natural Questions training set (Kwiatkowski et al., 2019)
TriviaQA training set (Joshi et al., 2017)
*
WebQuestions training set (Berant et al., 2013)
*
ARC-Easy and ARC-Challenge training sets (Clark et al., 2018)
Hellaswag training set (Zellers et al., 2019)
OpenBookQA training set (Mihaylov al., 2018)
PIQA training (Bisk et al., 2020)
SeiQ training set (Welbl et al., 2017)
SIQA training set (Sap et al., 2019)
Winogrande training set (Sakaguchi et al., 2021)
Our "math" training mixture that totals 66B tokens gathers:
55B tokens from Proof-Pile 2 (Azerbayev al., 2023) that contain AlgebraicStack (11B), Open WebMath (Paster et al., 2023) and ArXiv (29B).
2B tokens from OpenMathInstruuttIll we select the instances with a correct answer from the training set (Toshniwal al., 2024)
7B tokens from DeepMind math (Saxton et al., 2019)
2B tokens from the following instruction- datasets:
Math-Orea (Mitra al., 2024)
* TinyGSM (Liu et al., 2023) (we only select million examples from there).
* StackMathQA (Zhang, 2024)
* MAmmoTH2 (Yue et al., 2024) (we only select the mathstackexchange subset).
NuminaMath-CoT (NuminaMath, 2024) (duplicated times)
MetaMathQA (Yu et al., 2023) (duplicated 3 times)
18

---

Published as a conference paper at ICLR 2025
C ADDITIONAL EXPERIMENTS
In all our experiments in Section 5, we report the average accuracy performance obtained by our pre-trained models respectively world knowledge, commonsense and math benchmarks. Here, provide the results per task. In Subsection C.I, we display for each task, the downstream performance on per parameter basis (similar to Figure 1) and in Subsection C.2, we plot for each task, the downstream performance on a per validation perplexity basis (similar to Figure 6).
C.I DOWNSTREAM PERFORMANCE ON A PER PARAMETER BASIS
COMPLEXWEBQUESTIONS
HOTPOT-QA
NQ
16
F1 Accurddy (80
Accuracy (16)
18
Acouracy
12
15
31
16M
49M
159M
569M
2.18
16M
49M
159M
569M
2.18
16M
49M
159M
569M
2.18
Number total parameters
Number total parameters
Number of total parameters
TRIVIA-QA
WEBQUESTIONS
38
Acauracy
27
Accuracy
1.-
28
18
16M
49M
159M
569M
2,18
16M
<9M
1596
569M
2.18
Number parameters
Number
parameters
Dense transformer
Moli (18M active parametens)
MoE (58M active parameters)
Muli (20IM active parameters)
Figure 7: Downstream performance on the world knowledge tasks with respect to the total number of parameters of the models.
ARCC
ARCE
COMMONSENSEQA
37
64
43
Aecurscy
Accuracy
37
25
26
16M
150M
49M
159M
569M
2.18
16M
159M
569M
Number of total parameters
Number of total paramiters
Number of total parameters
HELLASWAG
OPENBOOKQA
PIQA
51
38
65
64
Accuracy (95)
47
Accuracy (%)
(96)
43
59
30
35
28
16M
49M
159M
569M
2.1B
16M
49M
159M
569M
2.18
15M
49M
159M
569M
218
Number of total parameters
Number of total parameters
Number of total parameters
SCIQ
SIQA
WINOGRANDE
40
Accuracy (16)
Accuracy (96)
45
Accuracy (36)
78
42
30
69
37
10M
49M
159M
509
2.18
10M
49M
159M
509/M
2.18
10M
49M
159M
506M
2.18
total porometers
Number
parameters
Number total parameters
Dense transformer
MoB (18M active parametens)
MoE (58M active parameters)
MoB (200M active parameters)
Figure& 8: Downstream performance on the commonsense tasks with respect to the total number of parameters of the models.
19

---

Published as a conference paper at ICLR 2025
GSM- HARD
GSMBK
HENDRYCKS MATH
45
49-
30
34
Accurasc [59
35-
(50)
Acarssy
19-
16-
10
16M
150M
569M
40M
159M
564M
2.18
16M
159M
569M
Number of total parameters
Number of total parameters
Number of total parameters
MINERVA
SVAMP
28
67
Accuray (9)
Accurary (3)
54
18-
14
16M
49M
159M
569M
2.18
16M
49M
159M
569M
2:18
Number total parsmesers
Number of total parameters
Dense transformer
MoE (ISM active parameters)
MoE (58M active parameters)
MolE (200M active parameters)
Figure 9: Downstream performance the math benchmarks with respect to the total number of
parameters of the models.
C.2 DOWNSTREAM PERFORMANCE ON PER VAL PERPLEXITY BASIS
COMPLEXWEBQUESTIONS
HOTPOT-QA
NQ
Accuracyy (8)
Accuracy
18
F1 Acturacy
12
19
14
7-
6.1
7.3
0.2
12.2
17.6
6.1
7.3
9.2
12.2
17.6
6.1
7.3
9.2
12.2
17.6
Validiatio perplexity
Validation perplesity
Validation perplexity
TRIVIA-Q QA
WEBQUESTIONS
38
Accuracy (85)
Accurryy (36)
25
27
18
18
7.3
12.2
17.5
6.1
7.3
9.2
12.2
17.6
Validation perplenity
Validation perplecity
Dense transformer
MoEi (1SM active parameters)
MoE (58M active parameters)
MoE (200M active purmmeeers)
Figure 10: Downstream performance on the world knowledge tasks with respect to the validation perplexity.
20

---

Published as a conference paper at ICLR 2025
ARCC
ARCE
COMMONSENSEQA
37
04-
48
Acturacy (50
Acturacy (50
57
49
37
40
26
34
17.5
6.1
0.2
12.2
17.6
Validation per perpiiiyy
Validation perplexity
Validation perplexity
HELLASWAG
OPENBOOKQA
PIQA
38-
65-
64
Accuray (9)
Accurary (9)
Accunny (9)
33
59
30
28
6.1
7.3
9.2
12.2
17.6
6.1
7.3
9.2
12.2
17.6
6.1
7.3
9.2
12.2
17.6
Validation per perplesity
Validation perplecity
Validation perplecity
SCIQ
SIQA
WINOGRANDE
48
Accuracy (95
Accurnzy (18)
45
Accurny (%5
78
42
39
68
37
6.1
7.3
9.2
12.2
17.5
6.1
7.3
9.2
12.2
17.6
6.1
7.3
9.2
12.2
17.6
Validation perplesity
Validation perplecity
Validation perplecity
Dense transformer
MoF: MoF active parameters)
MoF (58M active parametens)
Mof: (200M active parameters)
Figure 11: Performance on the commonsense tasks with respect to the validation perplexity.
GSM-HARD
GSMBK
MATH
45
49
30
Accuracy (X)
34-
Acturary (50)
35
Accurnny (X)
16
19
10
9-
3.2
3.5
3.9
4.5
5.5
3.2
3.5
3.9
4.5
3.2
3.5
3.9
4.5
5.5
Validation perplexity
Validation perplasity
Validation perplexity
MINERVA
SVAMP
28-
Accuracy (85)
Accorscy (98)
54
18
3.2
3.5
3.9
4.5
5.5
3.2
3.5
3.9
4.5
5.5
Validation parplaxity
Validation parplisity
Dense transformer
MoE (ISM active parametens)
MoE (58M active parameters)
Moli (200M active parameters)
Figure 12: Downstream performance on the math benchmarks with respect to the validation perplexity.
21

---

Published as a conference paper at ICLR 2025
D PROOFS
D.1 REASONING PROOFS
Definition D.1 (Set-disjointness task). Set disjointness the following task: given two inputs A, B E {0, 1)" for some r E N, compute max; B,
set-disjointness can be thought of as follows: Alice and Bob are given sets A and B respectively. Their objective to determine whether they have any overlapping items in their sets.
Lemma D.2 (Equivalence of set-disjointness and length-2 path). The ser-disjointmess task equiva- lent to the length-2 path task.
Proof. (=> ): Given an instance of set-disjointness, we can encode it into a length-2 path problem. Denote every item as vertex. Denote extra vertices as A, B, corresponding Alice and Bob. For every element i that Alice has, draw an edge between A and i. For every element that Bob has, draw an edge between B If and only if there are any overlapping elements, then there is a length-2 path from A to B. The number of elements because the number of vertices that do not belong to Alice Bob.
(=) ): Consider an instance G = (V, E), d of length-2 path, where s is the source vertex and d is the sink vertex. For all vertices with an edge with s, put this element into Alice's set of elements. For all vertices with an edge with d, put this element into Bobs's 's set of elements. If and only if there is length-2 path, then Alice and Bob's sets are overlapping. Then, is the number of vertices.
L
Lemma D.3 (Communication complexity lower-bound on concatenated outputs). For some sequence length, fix two disjoint subsets A, B C [N - 1], and consider a single-layer transformer f E m.n.1 with Ollog N)-bit precision that solves set disjointness for any input X where XA is function of Alice's 's input a E {0, 13", XB is function of Bob's imput b e {0, 1Y", and XINIMAUB) is fixed regardless of a, b. Then, f has width satisfying mH S2(r/log N).
Transformer N
Proof. By re-writing the following, the remainder of the proof from Sanford al. (2024) still holds. (
DISI(a, b) = vb ([softmax(6((a Qu K di (X))) Jen]nm))
This is because we may still use the same definition for Zn,s, Ln,s as in the proof. Hence, this concludes the proof.
D.1.1 PROOF OF THEOREM 3.2
We restate the corollary.
Theorem D.4 (Theorem 3.2). For some input sequence G = (V, fix two disjoint subsets A, B C
[N 1), and consider a single-layer transformer f E Transformer m. H.I.KKK with Ollog N) N precision that solves length-2 path for any input X where XA is a function of edges with the source X B is a function of edges with the destination d. Then, has width satisfying mH = 2((V)/ /og N).
Proof. The proof outline is follows:
Adapt Lemma 39 (Sanford et al., 2024) to support concatenation instead addition from different attention heads.
2. The lower bound with concatenation holds for length-2 path because set-disjointness and length-2 path are equivalent.
3. Extend the result to sparse transformers.
22

---

Published as a conference paper at ICLR 2025
We complete the first step with Lemma D.3. We complete the second set due to Lemma D.2. It remains to show that a router function also yields the same lower bound. We show that Lemma 39 of Sanford al. (2024) can be generalized to the case in which & is applied according routing function. Specifically, consider top-1 routing function r R + [K]. and K element-wise functions 19,,,... ,,,,.,, Rxx Rm R. For shorthand, define:
Y(Xx) [softmax (0(xN) (((txn)T@Qul "Q.K (()))
ME(M] J
which is the output of the attention head prior to applying the element transformation. Next, we define f(XN) ) as the output when the router function r used select 4);-
f(XN) = -I(r(((Xx)) = (Y
Because the lower bound does not place any restrictions on the function vi and rather argues a communication-copler. lower bound due information from Y(Xx), the lower bound also holds for routing function.
D.1.2 PROOF THEOREM 3.3
re-state Theorem 3.3 and give its proof.
Theorem D.5 (Theorem 3.3). For sequence length N, there exists f Transformer m,1,1 with O(log --u precision and width [V] that solves length-2 path for any input X.
Proof. Tokens are elements in V = U (0) x (0). The input is as follows: for vertex i, if the source shares an edge with that vertex, then the i'th input value is (s, i). Otherwise, it (s, 0). The first |V| tokens we correspond edges possibly shared with the source vertex. Then, the last |V] input tokens correspond to edges possibly shared with the destination vertex and share the same format as the first tokens. In hetween,
have arbitrary edges (u, v). We define embedding function where e, is the i'th standard basis vector in dimension r.
Φ : V RIVI
if 0 and u = 8 or ur if O.
vu,
Next, define V,, E | to be the identity matrix, and Qa, V/ VIXMV both have 0 everywhere. Consequently, the attention matrix is given by:
1/M :
1/
2
if there is path through i
((X)
1/IVI if one target vertex shares an edge with i
1/
1/V
0
otherwise.
For any entry that exceeds M the correct answer is there is a length-2 path. Hence, any thresholding function which achieves this separation suffices. Provided that any rounding T so that represented with Ollog N) is sufficient to separate between and 22, then O(log N) bits are sufficient.
D.2 MEMORIZATION PROOFS
Assume that K is a power of 2, and let m mo + (K). We associate each expert with a vector v; E [=1]))o(()) and choose r; (Vi,... (lo .,..., 0) E p
Lemma D.6. For some e > O, and some input (0, clm), the probability of routing to expert i 1/K K for all i.
Proof. By construction, we choose the expert as follows:
argmax U
jE[K]
The solution above must be the one whose signs match those of x. Because the entries of are O-mean, this implies that the probability of routing any particular expert is 1/K.
n
23

---

Published as a conference paper at ICLR 2025
Lemma D.7. Fix 8 E 10,1) and some expert j. With probability least 1 - 8. the number of examples routed to j is at most 2N K, given N 2
(l((1/8) examples from NO, clm).
2
Proof. The proof relies on a Hoeffding bound: for bounded independent random variables
X1,..., X such that a I X, I bii and SN = >i= X we have:
P[SN [[[ 2 E[SN]] S e 2 N(b a))
In our case, let X, 1 [sample i is routed expert j}. for some fixed [K]. Note that P[X;; 1] = 1/K due to Lemma D.6 and as such E[Sx] = N/K. Further, X, E [0, 1]. Hence we obtain:
2N K
2)N/K)²
P
SN Z
I exp
N
2.V
< exp
K2
In order to upper-bound this probability with 6, we obtain:
2N
8 M exp
K
2N
In
K2
K2 (n(6/))
< N.
2
Lemma D.8. Fix 8 E 1), K and the routing function as described in Lemma D.6. Given K² lm(1/8) samples with embedding dimension For every expert i, with probability at least
1 1- 8, there exists an MLP of width ((n/mK) that correctly classifies the examples routed to this expert.
Proof. We begin with result from Daniely (2020). Consider h, depth-2 network with neurons, an activation function : E + IR, which is OL))-Lipschitz, piecewise wwice-differentiable and satisfies E (x [x)o(XX] = O, weights E {+1} such that Σi 0i O(vg). Such a network,
defined by W and a, computes the function:
4 Σ a,o((w,, x)).
(w.a(x)
/q
For our proof consider ò be the absolute value function, and later show how to obtain MLP with ReLU activations.
Consider the set of samples ((x1, y1)...., (IN, (NN yN)).. )) where x; N/O,, Ims) and yi (+1) } is
Rademacher random variable. The objective is to find W a such that hihw- (2,) > Vi.
Mg Daniely (2020) gives the following result: for N < Tog+(M) and q 2 log*(M with probability 1 (1) there exists such that, for every E [N], yx(()) (In(M)).
We build on this result as follows: first, we assume that the attention matrix performs averaging of the input sequence, such that J; 인 ",",((,) where where X is the i-th input sequence to memorize. Assuming the inputs are also Gaussian, then so is the average. Hence, E R and I; ~ NO, cIm) for co 1
Next, we fix the number of experts to K. We fix router function r: Rlog K [K] such that, on average, each expert must memorize n/ K sequences, and with high probability, each expert must memorize at most 2n/ K sequences. We use the construction from Lemma D.6 where the router only acts the first log K entries of each vector. Hence, we have K MIPs, each tasked with memorizing
24

---

Published as a conference paper at ICLR 2025
at most imput-output values. We use the network construction from Daniely (2020) only the last m - log K 2 m/2 coordinates, which are not used by the router, so that their distribution is independent and Gaussian.
(m-log K)g Hence, with n S log (m- log K)) I
4 (m--(mo K)) and z log (m log K), with high probability, (x,)(z;)) = [[n(m))), where h2 is the expert to which is routed. It follows that width
4n log - log K)
mk
is sufficient for each individual expert, and we assume that 4n mk 2 1.
In order obtain MLP with ReLU activation defined (xx) - 0, 27). use the fact that [x| (z) + (-x)). This is because absolute value is a valid activation function in Daniely (2020), but ReLU is not. However, doubling the width of each MLP is sufficient obtain MLPs with ReLU activations instead, i.e.:
Sn log (m log K)
mK
This is done by taking solution which is of form:
kw,a(x) =
a [kw,,)))
and converting it to the form:
hw (2))
1 Σ (a(lo(w,, + (o(-w,, z))). V
The router has K vectors each of dimension m. Consequently, we need (KK)) parameters to store the router, Each expert has width 8n log"
K) Ö mkk and Ó k parameters.
mK
Corollary D.9. Let 8 E (0, 1), and fix K > For n 2 K² k21n(K88) with probability at least 1 -
there exists sparse transformer with K experts such that y;s(1;) Mlnm.). It has Ö(n + Km) parameters and Ó(n/ K + Km) active parameters.
Proof. For each expert, we apply the result from Lemma C.8 with 8' = 5 Hence, for every expert, with probability at most 5/ there does not exist an MLP of width Onn/m which memorizes its samples. By union bound, this implies that with probability at most ö, at least one of the experts is not able to memorize its samples. Hence, with probability at least 1 - 6, all of the experts are able to perform memorization on their given samples. In total, this implies we will use Ó + K)) parameters to store the entire mixture of expert network. The number of active parameters is Ö(n/ K + K). L
Lemma D.10 (Bound on norm of vector from N (0,cIm)) ). Let x ~ (0, cl for some c > O, and let 8 (0,1) 1). Then, with probability at least 1 8, there exists a constant Ca,m > 0, such that
(1)
2l (t)
||0||2 S Ca.a
+ 2 (( 1mm+2
Proof. Each component of x is distributed as x; ~ N'O, c). We can express as Ve2i, where 2 ~ (0, 1). Then,
|
1=1
The sum - 2? follows chi-squared distribution with m degrees of freedom. By the Laurent- Massart theorem Laurent & Massart (2000), for all > O, we have
ž m + 2/mm mt + 2t S Se
25

---

Published as a conference paper at ICLR 2025
Multiplying both sides inside the probability by obtain
P ( (m + 2/ml 21)) se
Setting / = In it follows that
+ 21n ( 11))
| 2 e m + 2 m In
s 6.
Therefore, with probability at least 1 - 5, we have
( ④
1 )
| S V m + 2 mln
21n
7
Lemma D.11 (Bounded bit complexity required). With high probability, (1) bits are required to store each weight the network, and Ö(n + K) bits are required to store the entire network. Ō n
K) active bits are required.
K
Proof. We show that the learning result from Daniely (2020) gives a network with bounded bit complexity. Hence, suffices to bound the bit complexity of the initial weights and the bit complexity of the gradient update. First, we begin with h(x) v E: ((w,, .²)). The objective is to show
there exists h(x) with bounded bit complexity which satisfies y;h(x;; > 0 vi. Suppose we begin by creating bins for the weights, and replace each w, with úū, where I|w u;lloo < e. Further, assume that || will oo < M for some M > 0. We then obtain that ||e Will2 < /mc. In addition, we state that with probability least 8, ||~||2 I Ca due to Lemma D.10. Using this, we have that:
[o<(w,, æ)) - ar ((Ø;, at)l (w; úi, .t)
is 1-Lipschitz) (Cauchy-Schwarr) (Lemma D.10).
I|w; wu @i||2||20|12
m w.p. 2 1 -
We apply this result to obtain that with probability at least 8,
1
[m(x)) - -(xx) <
Casllo((urrs - a((),x)))
/q
<TV/mccm mEca.m
Given that h(x)= 2 (1). then for some arbitrarily small constant (we use 4) we require: < V4V miv/mecc.m.....mm or equivalently, E S 4es.m Because h(#) 2 O(1,, we use this separation of to show that h(x) remains the same sign h(z) (however this constant can be replaced with an arbitrarily small constant so as to satisfy the requirement).
(0) Consider w to be the initialization of w; prior to the gradient step. Then, because, M = ||w lloc { (0)
we have that M max; { w. ah(x) = max; w; (0) log m @h(z) Dw:
{
}. Assuming that w is
Owi
initialized with bounded foo norm of 1, then obtain that M Lemma D.10 and that
C8,m. This is due to m//
log
dh(x))
a,00 ((w' (w),,,)))r;;
Ow;
if (w,, xx) < 0 otherwise.
1 0244
v9
26

---

Published as a conference paper at ICLR 2025
Using this, with high probability, require log( 22) log (8.Mss /qm) bits (by replacing E). Hence, with high probability, we require the number of bits to at most:
log
n
log m
cb.m )) - Ō(1)
We restate Theorem 3.6.
Theorem D.12 (Lower bound for dense model), Given the same task as above, a dense Transformer
requires (nn) parameters to solve the memorization task.
Proof of Theorem 3.6. Let e be the number of bits used for encoding each parameters (and we assume that is logarithmic in the problem parameters). Denote by H. the class of all transformers with W parameters and e bits per parameters. Since H is a finite class, where each function in the class can be encoded with eW bits, have (H] s 2-W Let X ..... XN E Raxd be the N input points. Assume a H. can solve the memorization task. Then, for every choice of y1...., UN E (+1), there exists transformer f E y1,... VN and therefore there are at least 2 different functions in H. So, we get 2 < [H] S 2cW and therefore W 2 N/c.
f(X,) = y; for all i [M.. There are 2 possible assignments for
Lemma D.13 (Active parameter comparison between dense and sparse transformers). There exist cases in which the amount of active parameters required to perform memorization is less for sparse transformer than dense transformer.
Proof. As shown, a dense transformer requires (1) parameters (and active parameters) to perform memorization of N sequences. In contrast, for sufficiently large and fixed K > 1, it holds that K K < n, which shows that the number of active parameters required is less for sparse transformer.
27


Reviews:
[{"review_id": "SMfBmycunH", "reviewer": "Reviewer_HxYz", "summary": "The paper explores the effectiveness of Mixture-of-Experts (MoE) architectures compared to standard dense Transformers, focusing on their abilities in memorization and reasoning tasks. Using both theoretical analysis and extensive experiments, the authors demonstrate that MoEs excel in tasks requiring high memorization capacity but fall short in reasoning tasks compared to dense models. Through synthetic tasks, such as shortest path and phone-book memorization, and real-world benchmarks in natural language and mathematics, the authors establish that while increasing the number of experts in MoEs enhances memorization, it does not lead to equivalent improvements in reasoning tasks.", "questions": "Generality of Findings: Could the authors discuss how their findings might generalize to larger models or different types of tasks, especially those that blend memorization and reasoning?\n\nRole of Chain of Thought Techniques: Considering that Chain of Thought prompting has been shown to enhance reasoning capabilities in language models by encouraging step-by-step reasoning, have the authors considered integrating CoT techniques into MoE architectures?  Chain of Thought techniques enable models to break down complex reasoning tasks into intermediate steps, which could help MoEs overcome the reasoning bottleneck identified in the paper. Exploring this integration might reveal new insights into how MoEs can be adapted or extended to better handle reasoning-intensive tasks. It would be interesting to investigate whether CoT could mitigate the limitations of MoEs on reasoning tasks and potentially improve their performance to match or exceed that of dense transformers.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "Theoretical Foundation: The paper offers theoretical analysis, including formal proofs and establishing clear upper and lower bounds on MoEs' capabilities for both memorization and reasoning tasks. Utilizes communication complexity theory to highlight fundamental differences in parameter scaling and capabilities between MoEs and dense Transformers.\n\nEmpirical Validation: Extensive experiments on both synthetic tasks (e.g., shortest path, phone-book memorization) and real-world benchmarks (natural language and mathematical reasoning tasks) support the theoretical claims.\n\nProvides insights that increasing the number of experts in MoEs enhances memorization performance but does not significantly improve reasoning abilities. Offers valuable guidance for model selection and scaling strategies, suggesting MoEs for memory-intensive tasks and dense models for reasoning tasks.", "weaknesses": "Limited Scale in Experiments: The largest models evaluated contain about 2.1 billion parameters, which is small compared to state-of-the-art models with tens or hundreds of billions of parameters. Uncertainty remains about how these results generalize to larger scales and whether larger MoEs might exhibit different performance characteristics.\n\nLimited Exploration of MoE Variants and Routing Strategies: Focuses mainly on standard MoE architectures with top-2 token-choice routing. Does not explore other routing mechanisms or architectural variations that might influence performance on reasoning tasks.\n\nDiversify Reasoning Tasks: The division of tasks into \"memorization\" and \"reasoning\" is somewhat binary. Many real-world tasks require a combination of both, which the paper does not deeply explore. It would be useful to include a broader range of reasoning tasks, such as logical reasoning, commonsense reasoning, multi-hop reasoning, and tasks from varied domains like code understanding. This could help analyze tasks that combine memorization and reasoning to reflect real-world complexities.\n\nAblation Studies: There is a lack of ablation studies on hyperparameters such as depth, width, and number of experts, which could provide more insights into the factors influencing performance.", "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": "fFPurw9tiz", "reviewer": "Reviewer_7TRg", "summary": "This paper explores the performance of the Mixture-of-Experts (MoE) architecture in tasks requiring different capabilities, specifically memorization and reasoning. It examines how increasing the number of experts in MoEs affects performance, particularly comparing MoEs to dense transformers. The authors find that while MoEs improve memorization substantially, they fall short in reasoning tasks compared to dense transformers. The paper includes both theoretical and empirical analyses across synthetic graph problems, closed-book retrieval tasks, and benchmarks in math and natural language, ultimately concluding that MoEs do well in memory-intensive tasks but struggle with reasoning.", "questions": "1. Fairness in Model Width Comparison: Why don’t the authors match the width of the MoE to that of the dense model? A wider MoE model could potentially invalidate some of the theoretical findings and should be considered for comparison.\n2. Practical Value of Theoretical Results: How does the theoretical analysis, particularly regarding graph problems, inform real-world applications? Are the assumptions made realistic? This should be stated explicitly\n3. Definition of Memorization: What is the precise definition of memorization used here? Typically, memorization implies reproducing the exact sequence encountered during training. Is this the criterion used? Or at least per token accuracy. \n3.1. Why do the authors think a large generalization gap is a good definition of memorization? Memorization should a statement only about the training performance. Once the test error is included we are instead thinking about lack of generalization, which is **not** the same as memorizing.\n4. Random Graph as a Reasoning Task: Why does the shortest path in random graphs qualify as reasoning? Tasks that can be solved by a deterministic algorithm do not generally require reasoning but rather algorithmic execution.\n5. Reverse Results in Practice: If the no free lunch theorem implies the necessity of specialized experts, then intuitively, MoEs should make sense for diverse tasks. Why, then, do the results show MoEs underperforming dense models in reasoning tasks? (note, the random graph problem has a deterministic algorithm and therefore does not count as real human reasoning. Coming up with the algorithm for that wold have been the reasoning)\n6. Memorization in Synthetic Tasks: For the phone book memorization task, how many epochs were used to achieve memorization? In other fields like vision, 100 epochs are often required for complete memorization. What was the exact train accuracy?\n7. Limitations in Problem Framing: Why not model MoEs with SFTed experts for specific domains (e.g., medical, math, translation) and then fine-tune the router? This approach seems more relevant in practice than the current setup, as it could allow routing to be specialized by domain without imposing limitations on MoEs’ capacity to compete with dense MLPs.\n8. Scaling Challenges in Reasoning: Why does scaling MoEs fail to match dense models in reasoning tasks? Can the theoretical results be explained intuitively? My intuition is that reasoning requires a lot of fact knowledge (eg Theorem Proving). Thus, shouldn’t “memorizing” models be able to retrieve facts better and reason about mathematics better? Perhaps this implies the benchmarks chosen weren’t good and should have chosen Olympiad Bench and the aforementioned MMLU benchmarks.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. Interesting: interesting question to differentiate the performance of MoEs and dense transformers across task types, particularly highlighting the unique strengths of MoEs vs MLPs in a task oriented way\n2. Clarity: The paper is clear and well-organized\n3. Glad generative tasks like GSM8K and MATH were included.\n4. Solid motivation of MoEs in the literature in intro.", "weaknesses": "Scientific Weakness\n1. Insufficient Empirical Validation: The paper would benefit from a broader set of empirical evaluations beyond the synthetic tasks used. For example, adding MMLU subtasks (eg Abstract Algebra and College mathematics are good starts. Under philosophy there is a formal logic task too. Big-Bench has an Induction task and Identify Theorem) and OlympiadBench [1] could provide richer insights into MoEs’ effectiveness on real-world tasks. This would also address the issue of over-reliance on benchmarks like GSM8K and MATH, which may be too narrow. Additional experiments at a larger model scale would also help validate the generalizability of the findings.\n2. Limited Real-World Applicability: Many of the toy experiments and theoretical results are not clearly connected to practical outcomes. This should be made directly and explicitly early in the paper. The assumptions made, especially regarding data distribution, need clearer justification to demonstrate their relevance in real-world applications. Without this, it’s difficult to assess if these experiments are accurate models for practical settings. A toy example that actually encompasses reasoning is needed for me to raise my score. Something that is real human reasoning and thus doesn’t have a deterministic algorithm that solves it (like random graphs or arithmetic).\n3. Alternative Approaches Left Unexplored: The paper does not examine whether MoEs would perform better if used with experts trained via supervised fine-tuning (SFT) on domain-specific data (e.g., separate experts for medical, math, coding, etc.). This approach may be more practical and representative of real-world applications, as it allows fine-tuning the routing based on relevant domain experts. The authors might believe this is a seemingly irrelevant point but part of the reviewing process is to identify if the authors solved the right problem. I am suggesting the paper did not solved the right problem and at least an argument of why their paper matters in practice compared to this suggested approach is important. Having multiple experts fine-tuned (or continually pre-trained) on high quantity on expert knowledge, combining them in a MoE then fine-tuning the router is likely a more realistic setting. If not, the authors should explain why not and more clearly why their results are relevant in practice. I’d like to know why the authors think they solved “the right problem”, see my question 7. \n4. Ambiguity in Memorization Definition: The authors equate a large train-test accuracy gap with “memorization,” but this is problematic. Memorization, requires the exact reproduction of a training target, whereas the LLM evaluations typically do not evaluate entire reasoning sequences, and train errors are not zero. The mathematical benchmarks used only check the final answer and for the MATH dataset use solvers like SymPy to report “reasoning accuracy”. This distinction in terminology affects the clarity and interpretation of the findings, and I would recommend revisiting and clarifying this.\n5. Question 1 in question section.\n6. Intuitive explanation is your results, see my question 8.\n7. Given my reservations, I think the title is overstated and requires a more scientific name and thus revision. I recommend removing the word parrot from it and change it in favor of something scientific.\n\nWilling to raise my score if authors address the weakness well, ideally with experiments. It’s already sorted in order of importance/urgency. \n\nWriting Weakness\n1. Abstract Needs Broader Implications: The final sentence of the abstract is vague about the implications of the work. Explicitly stating the broader significance of the findings.\n2. Lack of Connection to Main Contribution in Abstract: Phrases such as “we pre-train a series of MoEs and dense transformers and evaluate them on commonly used benchmarks in math and natural language” do not contribute to the clarity of the abstract. It would improve the paper if this evaluation connected back to the main contribution and quantified the practical significance of the results, which would help contextualize their importance.\n3. Overemphasis on Model Scale: The claim that language models’ capabilities stem from simply scaling model size (i.e., parameter count) is misleading. The quality, diversity, and alignment of data to task objectives (aligning with the “no free lunch” theorem) are equally critical in model performance. Thus, this paper underplays the impact of data performance, which should be crucially revised in the introduction.", "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": "SxEquurQFK", "reviewer": "Reviewer_A95z", "summary": "This paper investigates the difference between how performance scales with the number of experts on reasoning vs memorization tasks. Specifically, they show through theoretical results, synthetic datasets and common benchmarks that reasoning tasks scale more with the number of active parameters whereas memorization tasks scale more with the number of total parameters.", "questions": "Why did you choose different architectures for the synthetic and non-synthetic tasks?\n\nDid you consider looking at tasks that are a mix of both knowledge and reasoning (ex solving crossword puzzles)? It would be interesting to see more tradeoffs in these mixed tasks between scaling depth and number of experts.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "Tradeoffs between dense and sparse models is a very relevant problem as it’s a differentiator between many existing foundation models.\n\nThe authors clearly demonstrate an interesting finding that scaling experts improves knowledge/memorization tasks but has a lesser benefit on reasoning tasks (worse than dense activation-matched alternatives). \n\nThe results are very convincing, as this trend is demonstrated on theoretical tasks, synthetic tasks designed for reasoning/memorization and on popular benchmarks\n- Experiments in each section seem reasonable to back up the claim!\n\nGraphs are very clear and the overall story is very easy to follow.", "weaknesses": "More elaboration on what is meant by “this suggests that MoEs do have a bias towards learning functions that memorize the training data” when talking about figure 6a would be helpful:\n- Wouldn’t this bias show by just having lower perplexity? Why does the perplexity-accuracy relationship show this?\n\nMany tasks people care about require a mix of knowledge and reasoning. It would strengthen the study to do more targeted synthetic or real benchmarks designed to explicitly test the scaling properties for these tasks.", "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": "hPFVDYXULk", "reviewer": "Reviewer_toS5", "summary": "This paper proposes to study the performance tradeoffs between MoE models and standard dense transformers. More specifically, the paper proposes a theoretical analysis of MoE models at reasoning that validates the fact that MoEs can better handle knowledge-intensive tasks but are inferior to traditional dense transformers on those that need generalization ability for reasoning. Both synthetic datasets and real-world dataset verify the conclusion of this paper.", "questions": "I am curious to hear the authors thoughts regarding the reasoning ability of MoEs, if there exists possible future directions to mitigate the performance gap between MoEs and dense transformers? This is quite an important direction for both development of MoEs and LLM reasoning.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1. The paper is well-written and easy to follow, the problem this paper focuses is important to the community.\n2. Theoretical analysis is sound.\n3. The experiments are solid, the authors use synthetic data to provide a straightforward illustration and then extend to real-world datasets. The experiments support the claims of this paper sufficiently.", "weaknesses": "1. The proposed claim serves as a general conclusion. Have the authors tried on other reasoning tasks? For example, logical reasoning and code generation tasks are widely used to validate the effectiveness of reasoning. \n2. If the $k$ in top-k routing (k=2 in the main paper) influence the conclusion of this paper? And have the authors tried different routing machanisms?\n3. The analysis focuses on depth-1, however, in the synthetic data experiments, the authors use $L=12$ transformer layers to verify. Why does the authors not choose 1-layer transformer to verify the results on synthetic data?\n4. Compared to dense transformers, if there are configurations that can perform better than dense transformers theoretically?\n5. I'm curious about if the MoE performance on memory tasks and reasoning tasks is robust to the initialization or distribution of parameters.", "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 (Poster)