Mixture of Parrots: Experts improve memorization more than reasoning

## Abstract

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.

## 1 Introduction

The explosion in capabilities of large language models in recent years has largely been enabled by scaling their size, as 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 floating-point operations (FLOPs), making training and inference more computational intensive. Mixture-of-Experts (MoE) were introduced as a solution for this problem (Shazeer et al., ; Lepikhin et al., ; Fedus et al., ). 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.e., 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 on the MoE architecture (Achiam et al., ; Databricks, ; Anil et al., ; Dai et al., ; Jiang et al., ; Yang et al., ).

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-based 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-intensive tasks, we show that scaling the number of experts is competitive with scaling standard “dense” MLPs.

Figure: (a) Evaluation: world knowledge
Refer to caption: x1.png

To demonstrate these claims, we begin with a theoretical analysis of MoEs and dense models.
We use communication-complexity lower bounds to show that a single-layer MoE requires a 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 a 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 a 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 a dense transformer (4(b)).
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 (4(a)).

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-intensive tasks, MoEs surprisingly have a great advantage, where increasing the number of experts can match the performance of large dense models (1(a)). However, we show that for tasks that rely on reasoning, scaling the number of experts cannot compete with increasing the model dimension (Figures 1(b)-1(c)). Moreover, MoEs exhibit some memorization behaviors when trained on math problems (Figure 5). Taken together, our results show that the gains from using MoEs depend greatly on 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.

## 2 Related work

##### Mixture of Experts.

Mixture-of-Experts (MoE) date back to the work of Jacobs et al. (); Jordan & Jacobs (). Shazeer et al. (); Fedus et al. () were the first to scale this idea to deep learning and obtain state-of-the-art models in machine translation. Since then, several works have improved their routing algorithms (Lepikhin et al., ; Lewis et al., ; Roller et al., ; Clark et al., ; Zhou et al., ; Antoniak et al., ; Zhong et al., ), have improved their downstream performance after finetuning (Du et al., ; Zoph et al., ) or made their training and inference more efficient (Rajbhandari et al., ; Gale et al., ; Pan et al., ; Tan et al., ). However, only a few papers have studied the science of MoEs and their comparison with dense transformers. Clark et al. (); Krajewski et al. () establish scaling laws for MoEs. Chen et al. () design a specific classification problem where a model with multiple experts provably outperforms one with only one expert. Shazeer et al. (); Lepikhin et al. (); Artetxe et al. (); Lewis et al. (); Fedus et al. (); Du et al. () 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, we 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-sized MoEs perform as well as dense transformers.

##### Language models and memorization.

Large language models (LLMs) store a considerable amount of knowledge in their parameters (Petroni et al., ; Heinzerling & Inui, ).
They memorize useful knowledge such as facts and commonsense (Zhao et al., ).
Many works studied how memorization occurs in LLMs by developing tools to locate the knowledge in the model (Meng et al., ; Allen-Zhu & Li, ; Liu et al., ) or by tracking the training dynamics (Tirumala et al., ; Speicher et al., ). We draw inspiration from Allen-Zhu & Li () 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 are essentially question answering tasks on Wikipedia facts.
With respect to theoretical findings, Kim et al. (); Mahdavi et al. (); Madden et al. () provide upper bounds on the number of parameters needed for dense transformers to perform memorization tasks under various conditions.

##### Language models and reasoning.

In recent years, transformer-based language models have displayed remarkable effectiveness in solving a broad range of reasoning tasks. Specifically, the reasoning capabilities of transformers have been studied in the context of arithmetic problems (Jelassi et al., ; Cho et al., ; Hou et al., ; Zhou et al., ; McLeish et al., ; Lee et al., ), mathematical reasoning (Zhang et al., ; Imani et al., ; Wei et al., ) graph problems (Sanford et al., ; Fatemi et al., ; Jin et al., ; Wang et al., ) and code challenges (Shi et al., ; Zhu et al., ). Recently, state-of-the-art language models were used for solving complex math olympiad problems (DeepMind, ; NuminaMath, ; OpenAI, ).
With respect to theoretical findings, various works study the reasoning capabilities of transformers, relating their expressive power to other complexity classes and formal languages (Weiss et al., ; Zhou et al., ; Strobl et al., ). Other works study how chain-of-thought can improve the reasoning capabilities of language models in terms of expressive power and learnability (Abbe et al., ; Merrill & Sabharwal, ; Malach, ). However, the reasoning capabilities of MoE language models compared to their dense counterparts have received comparatively less attention.

## 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 MoEs can effectively memorize random inputs, requiring significantly less computational resources (active parameters) compared to dense models.

### 3.1 Setting

Consider a one-layer transformer $f\in\text{Transformer}_{m,H,1}^{N}$ which takes as input a sequence of length $N$ and has logarithmic bit-precision. $f$ embeds the input into dimension $m$ via the function $\phi$. $f$ has $H\geq 1$ attention heads, whose outputs are combined via concatenation before we apply point-wise function $\psi$ (^1^11In multi-layer Transformers, each layer outputs a vector of size $m$. However, since our focus in this section will be on binary classification problems, we will let the transformer output a single scalar, and we interpret the output of the final token as the prediction for the classification task.).

Define the parameters as $Q_{h},V_{h},K_{h}\in\mathbb{R}^{m\times m}$, $\phi:\mathcal{X}\rightarrow\mathbb{R}^{m}$, $\psi:\mathbb{R}^{m}\rightarrow\mathbb{R}$.
The output of $f$ is:

$$ $f({\bm{x}}_{1},\dots,{\bm{x}}_{N})=\psi\left(\left[\mathrm{softmax}\left(\phi( x_{N})^{\top}Q_{h}K_{h}^{\top}\phi(X)\right)\phi(X)V_{h}\right]_{h\in[H]} \right).$ $$

$f$ is a *dense* transformer, if $\psi$ is an MLP, i.e. function of the form:

$$ $\psi({\bm{x}})={\bm{u}}^{\top}\sigma({\bm{W}}{\bm{x}}+{\bm{b}})\mathrm{,~{}~{} for~{}}{\bm{W}}\in\mathbb{R}^{m^{\prime}\times m},b\in\mathbb{R}^{m^{\prime}}, {\bm{u}}\in\mathbb{R}^{m^{\prime}}$ $$

where $\sigma$ is the ReLU activation function.
$f\in\text{Transformer}_{m,H,1,K}^{N}$ is a *sparse* (MoE) transformer with $K$ experts if $\psi$ is a function of the form:

$$ $\psi({\bm{x}})={\mathbf{u}}_{i}^{\top}\sigma({\bm{W}}_{i}{\bm{x}}+{\bm{b}}_{i} )\mathrm{~{}for~{}}i=\operatorname*{\rm argmax}_{j}{\mathbf{r}}_{j}^{\top}{\bm {x}}$ $$

where ${\bm{W}}_{1},\dots,{\bm{W}}_{k}\in\mathbb{R}^{m^{\prime}\times m}$, ${\bm{b}}_{1},\dots,{\bm{b}}_{k}\in\mathbb{R}^{m^{\prime}}$, ${\mathbf{u}}_{1},\dots,{\mathbf{u}}_{k}\in\mathbb{R}^{m^{\prime}}$ are the parameters of each expert and $r_{1},\dots,r_{k}$ define the routing function (we use top-1 routing).

### 3.2 MoEs require a critical hidden size to solve graph reasoning tasks

In this section, we analyze the graph reasoning capabilities of dense and sparse transformers. We define the length-2 path problem on a graph, and use it as a means to understand other graph reasoning tasks such as graph connectivity, shortest path, and cycle detection.

###### Definition 3.1 (Length-2 Path Problem) .

The input is a graph $G=(V,E)$. The source $s\in V$ and a destination $d\in V$ are fixed for all tasks as the 0 0 and $\lvert V\rvert$ vertex. The length-2 path problem asks whether there is a path of length $2$ from $s$ to $d$.

Graph connectivity, shortest path, and cycle detection are all graph reasoning tasks which reduce to the length-2 path problem due to (Sanford et al., ) and Lemma C.2.
We provide a lower-bound on the width required for a sparse transformer to solve the length-2 path problem, and an upper-bound on the width required for a dense transformer to solve the problem. Further, we show a separation 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 a dense model can solve with the same amount of *total* parameters.

###### Definition 3.1 (Length-2 Path Problem) .

##### 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 cycle detection.

###### Theorem 3.2 (Length-2 path lower-bound on sparse transformers) .

For some input sequence $G=(V,E)$, fix two disjoint subsets $A,B\subset[N-1]$, and consider a single-layer transformer $f\in\text{Transformer}_{m,H,1,K}^{N}$ with $O(\log N)$-bit precision that solves length-2 path for any input $X$ where $X_{A}$ is a function of edges with the source $s$, $X_{B}$ is a function of edges with the destination $d$. Then, $f$ has width satisfying $mH=\Omega(\lvert V\rvert/\log N)$.

The proof follows almost identically from the proof in (Sanford et al., ) for the class $\text{Transformer}_{m,H,1}^{N}$. The original proof does not place constraints on the function $\psi$ and is based on a communication-complexity argument. As such we may design $\psi$ so that it first routes and then chooses which expert to apply. We give a complete proof in Appendix C. As such, the result of (Sanford et al., ) can also be extended to the class $\text{Transformer}_{m,H,1,K}^{N}$.

###### Theorem 3.2 (Length-2 path lower-bound on sparse transformers) .

##### Upper bound on width of depth-1 dense transformer for reasoning.

In this section we give an upper bound for the width required for a dense model to solve the length-2 path problem.

###### Theorem 3.3 (Length-2 path width upper bound for transformer) .

There exists a transformer of width $\lvert V\rvert$, $H=1$, and $O(\log N)$-bit 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 $v$, the source and destination vertices, have edges with a common vertex. We defer the proof to Appendix C.

###### Theorem 3.3 (Length-2 path width upper bound for transformer) .

##### Parameter-matched comparison of dense and sparse depth-1 transformers.

Using the lower-bound on width required for a 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)^{2}$, the sparse model is unable to solve the same task as the dense model.

###### Corollary 3.4 .

Consider a sparse transformer (with $K$ experts) and a dense transformer with the same number of parameters. There exists a 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 a mixture of experts with $K$ experts. Let the width of the dense model be $m_{d}$, and the width of the sparse model be $m_{s}$. The number of parameters in the dense model is $O(m_{d}^{2})$ and the number of parameters in the sparse model is $O(Km_{s}^{2})$. In order to match the number of parameters, it must be the case that $m_{s}=\frac{m_{d}}{\sqrt{K}}$.
Suppose we let $m_{d}=\lvert V\rvert$, as this is sufficient to solve the above problems. For any $K\geq\Omega\left((\log N)^{2}\right)$, the sparse model is not sufficiently wide to solve the problem.
∎

###### Corollary 3.4 .

###### Proof.

### 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 a sparse transformer 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 of bits to encode each parameter.

We consider sequences $\{(X^{i},y_{i})\}_{i=1}^{n}$ where $X^{i}\in\mathbb{R}^{N\times m}$ are input sequences of length $N$ in dimension $m$ such that $X^{i}[j]$ is sampled from a Gaussian distribution $\mathcal{N}(0,I_{m})$. We assume $y_{1},\dots,y_{N}\in\{\pm\bm{1}\}$ are arbitrary labels for the $n$ sequences. The objective is for a transformer to memorize these sequences, i.e. map each input $X^{i}$ to a label $y_{i}$. The classification is determined by the sign of the last token output.

##### Upper-bound on MoE for memorization.

We begin by showing that, with high probability over the choice of the inputs, the MoE architecture can memorize (i.e., arbitrarily label the examples), with a 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 $\tilde{O}\left(\dfrac{n}{K}+Km\right)$ active parameters and $\tilde{O}\left(n+Km\right)$ total parameters stored in $\tilde{O}(1)$ bits that, when applied to each sequence $X^{i}$, outputs at the last token a value whose sign matches $y_{i}$, i.e.,
$\operatorname{sign}(f(X_{i}))=y_{i}\quad\text{for all }i=1,\dots,n.$(^2^22We use $\tilde{O}$ and $\tilde{\Omega}$ to hide logarithmic factors.)

Specifically, if we choose $K=\sqrt{n/m}$ we get that an MoE architecture can solve the memorization problem with $\tilde{O}(\sqrt{nm})$ 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 $\tilde{O}(n/{mK})$ neurons can memorize a sample of size $O(n/K)$. We present the proof in Appendix C.2.

###### Theorem 3.5 .

##### Lower bound on memorization with dense Transformer.

Next, we give a lower-bound on the number of parameters for a dense transformer to perform memorization.

###### Theorem 3.6 (Lower bound for dense model) .

Given the same task as above, a dense Transformer requires $\tilde{\Omega}(n)$ parameters to solve the memorization task.

This bound follows from the fact that there are $2^{n}$ possible labels for any fixed set of $n$ inputs, and at most $2^{cW}$ functions with $W$ parameters and $c$ bit per parameters. The proof is in Appendix C.2.

###### Theorem 3.6 (Lower bound for dense model) .

##### 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 $\tilde{O}(\sqrt{nm})$ active parameters can memorize, while a dense model requires $\tilde{\Omega}(n)$ parameters. So, for $n\gg m$, MoEs are significantly more efficient. Comparing the number of total parameters, MoEs require $\tilde{O}(n+Km)$ parameters, so both MoE and dense models have linear dependence on $n$ 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 a 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 memorization ability of our models. Our experiments in this section highlight that adding experts yields greater performance improvements on memorization tasks than reasoning tasks.

Figure: Figure 2: Illustration of the shortest path task. We feed the model with a sequence that lists all the edges in the input graph and ends with the query (in green) which asks the model to find a shortest path between two vertices (from vertex 1 to vertex 4 in the figure). The model then autoregressively returns the shortest path (in purple).
Refer to caption: x4.png

### 4.1 Experimental setup

##### Architecture.

We opt for the Mistral (Jiang et al., ) and Mixtral (Jiang et al., ) 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, we fix the number of layers to $L=12$. For the dense transformers, we vary model size by sweeping the width $d\in\{256,512,1024\}$. For MoEs, we sweep over widths $d\in\{256,512\}$ and the number of experts $E\in\{8,16,32,64\}$. To be consistent with our experiments in Section 5, we set the intermediate dimension in the FFN block to be equal to $d$ (and not $4d$). We use token-choice routing, do not apply any token dropping and each token is routed to 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 as the total number of parameters.

##### Shortest path task.

For a graph with $n$ vertices, our token space $\mathcal{V}$ is of size $n+6$ with tokens encoding the vertices and some special tokens: $\mathcal{V}=\{1,\dots,n,\left\langle{\textsc{EDGE}}\right\rangle,\left\langle{
\textsc{BOS}}\right\rangle,\left\langle{\textsc{EOS}}\right\rangle,\left
\langle{\textsc{PAD}}\right\rangle,\left\langle{\textsc{SEP}}\right\rangle,/\}$ where $\left\langle{\textsc{BOS}}\right\rangle$ is the beginning of sentence token, $\left\langle{\textsc{EOS}}\right\rangle$ the end of sentence token, $\left\langle{\textsc{PAD}}\right\rangle$ the padding token, $\left\langle{\textsc{EDGE}}\right\rangle$ is the token indicating an edge between two vertices and, $\left\langle{\textsc{SEP}}\right\rangle$ 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 Erdös-Rényi model, with $n$ vertices and probability $p$ for each edge to exist. We vary $n\in\{25,30,50,40,45,50,55\}$ and set $p$ such that the average length of the shortest path is 3.5. Each train/test pair corresponds to *one* value of $(n,p)$, we do not mix graph sizes.

Figure: 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, we randomly select a name in the phone-book (in green) and ask the model to return their phone number (in purple) without access to the phone-book.
Refer to caption: x6.png

Figure: (a) Phone-book memorization
Refer to caption: x7.png

##### Phone-book task.

Our token space $\mathcal{V}$ is of size 39 and made of the alphabet letters, digits and special tokens: $\mathcal{V}=\{a,\dots,z,0,\dots,9,\left\langle{\textsc{BOS}}\right\rangle,
\left\langle{\textsc{EOS}}\right\rangle,\left\langle{\textsc{SEP}}\right\rangle\}$. We generate phone-books where the names consist of 5 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 $1\mathrm{e}6$ and the test set consists of $1\mathrm{e}3$ held-out examples that are sampled from the *same* distribution as the training examples. For the phone-book experiments, we vary the training set size over $\{1\mathrm{e}{5},5\mathrm{e}{5},1\mathrm{e}{6},1.5\mathrm{e}{6},2\mathrm{e}{6}
,2.5\mathrm{e}{6},3\mathrm{e}{6}\}$ and the test set consists of $1\mathrm{e}3$ queries from the training set.

##### Optimization.

We use the AdamW optimizer (Loshchilov et al., ) with a weight decay equal to 0.1. We sweep the learning rate over $\{5\mathrm{e}{-5},1\mathrm{e}{-4},5\mathrm{e}{-4},1\mathrm{e}{-3}\}$, the number of epochs over $\{2,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 a linear decay scheduler. All models are trained by next-token prediction. In the graph task, we apply a mask on the input instance so that we only penalize the model whenever it makes a mistake on the labels (and not on the inputs and labels jointly). In the phone-book experiment, we do not apply any masking.

##### Evaluation.

For each task we compute the exact accuracy, i.e. we count the generation as correct only if it fully matches the ground truth. For the phone-book task, we report the size of 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 of different sizes and at test time, evaluate whether they memorize the phone number of some names.
4(a) 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.

The findings are 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 are only routing to 2 out of 64 experts. For instance, we find that an 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 a 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. 4(b) 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 a 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 corroborate 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, we will further corroborate this idea through evaluation of pre-trained models on commonsense and math reasoning benchmarks.

## 5 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.
The memorization-intensive tasks test for “world knowledge” and consist of benchmarks like TriviaQA (Joshi et al., ).
We break the reasoning-intensive tasks into two subcategories: one for natural language reasoning tasks like WinoGrande (Sakaguchi et al., ) and another for mathematical reasoning tasks like Hendrycks-MATH (Hendrycks et al., ).
These tasks may be seen as real-world analogs of the stylized phone-book and shortest path tasks studied in Section 4.

We observe that performance on world-knowledge tasks is governed by the total number of parameters 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 at 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).

### 5.1 Setup

##### Architecture.

We train dense transformers and MoEs using the OLMoE
codebase (Muennighoff et al., ). We set the number of layers $L=20$ and vary the width $d\in\{256,512,1024,2048,4096\}$ for dense transformers and $d\in\{256,512,1024\}$ for MoEs. Similarly to Muennighoff et al. (), 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\in\{8,16,32,64\}$. For the specific case of width 256, we
also train a MoE with 256 experts because its parameter count approximately matches the one of a width-2048 dense model and thus, we can compare the downstream performance of the two models. We use top-2 token-choice routing, without token dropping which is implemented in the dMoE function from the Megablocks package (Gale et al., ).

##### Training hyperparameters.

We use the AdamW optimizer (Loshchilov et al., ) with a 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 a linear decay scheduler. We train our models using FSDP (Zhao et al., ).

##### Pre-training datasets.

We train two collections of models, one collection on natural language and another one on math. The natural language dataset is a mixture constituted of FineWeb-edu (Penedo et al., ), Cosmopedia (Ben Allal et al., ), 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., ) and instruction datasets such as OpenMathInstruct (Toshniwal et al., ) and MetaMathQA (Yu et al., ). Each of the two training mixture approximately totals 65B tokens. A precise description of the training mixtures can be found in Appendix A.

##### Evaluation.

We measure the validation perplexity on 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., ), Natural Questions (Kwiatkowski et al., ), HotpotQA (Yang et al., ), WebQuestions (Berant et al., ), ComplexWebQuestions (Talmor & Berant, ).
- –
Commonsense tasks: ARC-C and ARC-E (Clark et al., ), CommonsenseQA (Talmor et al., ), HellaSwag (Zellers et al., ), OpenbookQA (Mihaylov et al., ), PIQA (Bisk et al., ), SciQ (Welbl et al., ), SIQA (Sap et al., ), WinoGrande (Sakaguchi et al., ).
- –
Math benchmarks: SVAMP (Patel et al., ), GSM8k (Cobbe et al., ), GSM-Hard (Gao et al., ), Hendrycks-MATH (Hendrycks et al., ) and Minerva-MATH (Lewkowycz et al., ).

Figure: (a)
Refer to caption: x9.png

In all our experiments, we plot the average accuracy for each of these three categories. We report the corresponding per-task performance in Appendix B.

### 5.2 Results

##### Experts improve memorization more than reasoning.

We observe that the conclusions from our theoretical results and synthetic experiments also hold when pre-training and evaluating language models on natural language and math. In 1(a), we report the accuracy of our models with respect 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 commonsense and math benchmarks (Figures 1(b),1(c)) we find that MoEs do not reach the performance of dense models with the same number of total parameters. This indicates that for these reasoning tasks, increasing the dense model width is more effective that adding experts.

##### 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.
In Figure 5 we select 6,319 random problems from the OpenMathInstruct dataset, which is part of the training mixture data. More precisely, we pick 5,000 Hendrycks-MATH 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 OpenMathInstruct dataset, Figure 5 shows that at scales beyond 159M parameters, MoEs suffer from a more significant generalization gap than dense transformers.
This is suggestive that MoEs are more prone to memorize training data than dense models.

##### MoE models excel at world knowledge tasks but match dense models in reasoning when perplexity 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 a particular model class is more likely to learn. 6(a) shows that at a fixed perplexity, the MoE models outperform the dense models on world knowledge tasks. This suggests that MoEs do have a bias towards learning functions that memorize training data. On the other hand, Figures 6(b) and 6(c) show that MoEs and dense models perform about the same on the reasoning tasks at fixed validation perplexity.
We can square this with the results from Figure 1 by noting that at equal total number of parameters an MoE has worse validation perplexity than the corresponding dense model. This suggests that while MoEs do not change the relationship between perplexity and downstream accuracy on reasoning tasks relative to dense models, they may struggle to learn the reasoning parts of the training distribution as well.

Overall, our main findings in Figure 1 and supplementary experiments in Figures 5 and 6 corroborate the hypothesis that MoEs can effectively use more experts to increase their memory capacity, but not necessarily their capability to reason.

Figure: (a)
Refer to caption: x11.png

## 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 $L$ has number of parameters which scales with $O(d^{2})$, and run-time that scales with $O(d^{2}L^{2})$.
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 et al., ; Peng et al., ; Fu et al., ; Gu & Dao, ), while MoEs attempt to improve dependence on $d$ 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-intensive tasks increasing the dimension $d$ is inevitable, and scaling the computation with $O(d^{2})$ seems unavoidable.
This remains true regardless of the different design choices in the MoE architecture and is backed up empirically.
We note that there is increasing interest in developing non-MoE models with sub-quadratic dependence on $d$, using some structural assumptions on the weight layers (Kamalakara et al., ; Dao et al., ; ; Fu et al., ), 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 reasoning-intensive tasks.

##### Limitations and future work.

While we provide substantial experiments on pre-trained models with up to $\leq 2.1$B parameters, we recognize that large scale MoEs like Mixtral (Jiang et al., ), DeepSeek-V2 (Dai et al., ), 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, it would be 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.

## 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-2377, and the
National Science Foundation Grant under award #IIS 2229881.

## Appendix A Details on the pre-training datasets

In Section 5, we pretrain two collections of models, one on “natural language” and the other on “math”. Here, we give a 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 et al., ).
- –
14B tokens from Cosmopedia (Ben Allal et al., ).
- –
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 to increase their presence in the mixture. The presence of these datasets is pretty important as argued in Allen-Zhu & Li () so that the model is familiar with the downstream tasks at test time.
$\ast$
ComplexWebQuestions training set (Talmor & Berant, )
$\ast$
HotPotQA training set (Yang et al., )
$\ast$
Natural Questions training set (Kwiatkowski et al., )
$\ast$
TriviaQA training set (Joshi et al., )
$\ast$
WebQuestions training set (Berant et al., )
$\ast$
ARC-Easy and ARC-Challenge training sets (Clark et al., )
$\ast$
Hellaswag training set (Zellers et al., )
$\ast$
OpenBookQA training set (Mihaylov et al., )
$\ast$
PIQA training set (Bisk et al., )
$\ast$
SciQ training set (Welbl et al., )
$\ast$
SIQA training set (Sap et al., )
$\ast$
Winogrande training set (Sakaguchi et al., )

Our “math” training mixture that totals 66B tokens gathers:

- –
55B tokens from Proof-Pile 2 (Azerbayev et al., ) that contain AlgebraicStack (11B), OpenWebMath (Paster et al., ) and ArXiv (29B).
- –
2B tokens from OpenMathInstruct-1: we select the instances with a correct answer from the training set (Toshniwal et al., )
- –
7B tokens from DeepMind math (Saxton et al., )
- –
2B tokens from the following instruction-like datasets:
$\ast$
Math-Orca (Mitra et al., )
$\ast$
TinyGSM (Liu et al., ) (we only select 1 million examples from there).
$\ast$
StackMathQA (Zhang, )
$\ast$
MAmmoTH2 (Yue et al., ) (we only select the mathstackexchange subset).
$\ast$
NuminaMath-CoT (NuminaMath, ) (duplicated 3 times)
$\ast$
MetaMathQA (Yu et al., ) (duplicated 3 times)

## Appendix B Additional experiments

In all our experiments in Section 5, we report the average accuracy performance obtained by our pre-trained models on respectively world knowledge, commonsense and math benchmarks. Here, we provide the results per task. In Subsection B.1, we display for each task, the downstream performance on a per parameter basis (similar to Figure 1) and in Subsection B.2, we plot for each task, the downstream performance on a per validation perplexity basis (similar to Figure 6).

### B.1 Downstream performance on a per parameter basis

Figure: Figure 7: Downstream performance on the world knowledge tasks with respect to the total number of parameters of the models.
[Uncaptioned image]: x14.png

Figure: Figure 8: Downstream performance on the commonsense tasks with respect to the total number of parameters of the models.
Refer to caption: x19.png

Figure: Figure 9: Downstream performance on the math benchmarks with respect to the total number of parameters of the models.
Refer to caption: x28.png

### B.2 Downstream performance on a per val perplexity basis

Figure: Figure 10: Downstream performance on the world knowledge tasks with respect to the validation perplexity.
Refer to caption: x33.png

Figure: Figure 11: Performance on the commonsense tasks with respect to the validation perplexity.
Refer to caption: x38.png

Figure: Figure 12: Downstream performance on the math benchmarks with respect to the validation perplexity.
Refer to caption: x47.png

blabla

blabla

## Appendix C Proofs

### C.1 Reasoning proofs

###### Definition C.1 (Set-disjointness task) .

Set disjointness is the following task: given two inputs $A,B\in\{0,1\}^{r}$ for some $r\in\mathbb{N}$, compute $\max_{i}A_{i}B_{i}$.

Set-disjointness can be thought of as follows: Alice and Bob are given sets $A$ and $B$ respectively. Their objective is to determine whether they have any overlapping items in their sets.

###### Lemma C.2 (Equivalence of set-disjointness and length-2 path) .

The set-disjointness task is equivalent to the length-2 path task.

###### Proof.

$(\implies)$: Given an instance of set-disjointness, we can encode it into a length-2 path problem. Denote every item $i$ as a vertex. Denote two extra vertices as $A$, $B$, corresponding to Alice and Bob. For every element $i$ that Alice has, draw an edge between $A$ and $i$. For every element $i$ that Bob has, draw an edge between $B$ to $i$. 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 or Bob.

$(\impliedby)$: Consider an instance $G=(V,E),s,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 set of elements. If and only if there is a length-2 path, then Alice and Bob’s sets are overlapping. Then, $r$ is the number of vertices.
∎

###### Lemma C.3 (Communication complexity lower-bound on concatenated outputs) .

For some sequence length, fix two disjoint subsets $A,B\subset[N-1]$, and consider a single-layer transformer $f\in\text{Transformer}_{m,H,1}^{N}$ with $O(\log N)$-bit precision that solves set disjointness for any input $X$ where $X_{A}$ is a function of Alice’s input $a\in\{0,1\}^{r}$, $X_{B}$ is a function of Bob’s input $b\in\{0,1\}^{r}$, and $X_{[N]\setminus(A\cup B)}$ is fixed regardless of $a,b$. Then, $f$ has width satisfying $mH=\Omega(r/\log N)$.

###### Proof.

By re-writing the following, the remainder of the proof from Sanford et al. () still holds.

$$ $\text{DISJ}(a,b)=\psi\left(\left[\mathrm{softmax}\left(\phi(x_{N})^{\top}Q_{h} K_{h}^{\top}\phi(X)\right)\phi(X)v_{h}\right]_{h\in[H]}\right).$ $$

This is because we may still use the same definition for $Z_{h,S},L_{h,S}$ as in the proof.
Hence, this concludes the proof.
∎

###### Definition C.1 (Set-disjointness task) .

###### Lemma C.2 (Equivalence of set-disjointness and length-2 path) .

###### Proof.

###### Lemma C.3 (Communication complexity lower-bound on concatenated outputs) .

###### Proof.

#### C.1.1 Proof of Theorem 3.2

We restate the corollary.

###### Theorem C.4 (Theorem 3.2 ) .

For some input sequence $G=(V,E)$, fix two disjoint subsets $A,B\subset[N-1]$, and consider a single-layer transformer $f\in\text{Transformer}_{m,H,1,K}^{N}$ with $O(\log N)$-bit precision that solves length-2 path for any input $X$ where $X_{A}$ is a function of edges with the source $s$, $X_{B}$ is a function of edges with the destination $d$. Then, $f$ has width satisfying $mH=\Omega(\lvert V\rvert/\log N)$.

###### Proof.

The proof outline is as follows:

- 1.
Adapt Lemma 39 (Sanford et al., ) to support concatenation instead of 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.

We complete the first step with Lemma C.3. We complete the second set due to Lemma C.2. It remains to show that a router function also yields the same lower bound.
We show that Lemma 39 of Sanford et al. () can be generalized to the case in which $\psi$ is applied according to a routing function. Specifically, consider a top-$1$ routing function $r:\mathbb{R}^{m}\rightarrow[K]$, and $K$ element-wise functions $\psi_{1},\ldots,\psi_{K}:\mathbb{R}^{m}\rightarrow\mathbb{R}$. For shorthand, define:

$$ $Y(X_{N})=\left[\mathrm{softmax}\left(\phi(x_{N})^{\top}Q_{h}K_{h}^{\top}\phi(X )\right)\phi(X)v_{h}\right]_{h\in[H]},$ $$

which is the output of the attention head prior to applying the element-wise transformation.
Next, we define $f(X_{N})$ as the output when the router function $r$ is used to select $\psi_{i}$.

$$ $f(X_{N})=\sum_{i\in K}{\mathbf{I}}\left\{r(Y(X_{N}))=i\right\}\psi_{i}(Y(X_{N} )).$ $$

Because the lower bound does not place any restrictions on the function $\psi$ and rather argues a communication-complexity lower bound due to information from $Y(X_{N})$, the lower bound also holds for a routing function.
∎

###### Theorem C.4 (Theorem 3.2 ) .

###### Proof.

#### C.1.2 Proof of Theorem 3.3

We re-state Theorem 3.3 and give its proof.

###### Theorem C.5 (Theorem 3.3 ) .

For sequence length $N$, there exists $f\in\text{Transformer}_{m,1,1}^{N}$ with $O(\log N)$-bit precision and width $\lvert V\rvert$ that solves length-2 path for any input $X$.

###### Proof.

Tokens are elements in $\mathcal{V}=V\cup\{0\}\times V\cup\{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 is $(s,0)$. The first $\lvert V\rvert$ tokens we see correspond to edges possibly shared with the source vertex. Then, the last $\lvert V\rvert$ input tokens correspond to edges possibly shared with the destination vertex and share the same format as the first $r$ tokens. In between, we can have arbitrary edges $(u,v)$.
We define an embedding function where $\mathbf{e}_{i}$ is the $i$’th standard basis vector in dimension $r$.

$$ $\displaystyle\phi:\$ $\displaystyle\mathcal{V}\rightarrow\mathbb{R}^{\lvert V\rvert}$ $\displaystyle(u,v)\mapsto\begin{cases}\mathbf{e}_{i}&\text{if }i>0\text{ and } u=s\text{ or }u=v\\ \mathbf{0}&\text{if }i=0.\end{cases}$ $$

Next, we define $V_{h}\in\mathbb{R}^{\lvert V\rvert\times\lvert V\rvert}$ to be the identity matrix, and $Q_{h},V_{h}\in\mathbb{R}^{\lvert V\rvert\times\lvert V\rvert}$ both to have 0 0 everywhere. Consequently, the attention matrix is given by:

$$ $\left(\begin{bmatrix}1/\lvert V\rvert&\ldots&1/\lvert V\rvert\\ \vdots&\ddots\\ 1/\lvert V\rvert&&1/\lvert V\rvert\end{bmatrix}\phi(X)\right)_{j,i}=\begin{ cases}2/\lvert V\rvert&\text{if there is a path through $i$}\\ 1/\lvert V\rvert&\text{if one target vertex shares an edge with }i\\ 0&\text{otherwise.}\end{cases}$ $$

For any entry that exceeds $\frac{1}{\lvert V\rvert}$, the correct answer is there is a length-2 path. Hence, any thresholding function which achieves this separation suffices. Provided that any rounding $\frac{1}{\lvert V\rvert}$ to $\tilde{c}$ so that it is represented with $O(\log N)$ is sufficient to separate between $\tilde{c}$ and $2\tilde{c}$, then $O(\log N)$ bits are sufficient.
∎

###### Theorem C.5 (Theorem 3.3 ) .

###### Proof.

### C.2 Memorization Proofs

Assume that $K$ is a power of 2, and let $m=m_{0}+\log(K)$. We associate each expert $i$ with a vector ${\mathbf{v}}_{i}\in\{\pm 1\}^{\log(K)}$ and choose ${\mathbf{r}}_{i}=(v_{1},\dots,v_{\log(K)},0,\dots,0)\in\mathbb{R}^{m}$.

###### Lemma C.6 .

For some $c>0$, and some input $x\sim\mathcal{N}(0,cI_{m})$, the probability of routing to expert $i$ is $1/K$ for all $i$.

###### Proof.

By construction, we choose the expert as follows:

$$ $i=\operatorname*{\rm argmax}_{j\in[K]}v_{j}^{\top}x.$ $$

The solution above must be the one whose signs match those of $x$. Because the entries of $x$ are 0 0-mean, this implies that the probability of routing to any particular expert is $1/K$.
∎

###### Lemma C.7 .

Fix $\delta\in(0,1)$ and some expert $j$. With probability at least $1-\delta$, the number of examples routed to $j$ is at most $2N/K$, given $N\geq\frac{K^{2}\ln\left(1/\delta\right)}{2}$ examples from $\mathcal{N}(0,cI_{m})$.

###### Proof.

The proof relies on a Hoeffding bound: for bounded independent random variables $X_{1},\ldots,X_{N}$ such that $a\leq X_{i}\leq b\ \forall i$, and $S_{N}=\sum_{i=1}^{N}X_{i}$, we have:

$$ $\operatorname*{\mathbb{P}}[S_{N}-\operatorname*{\mathbb{E}}[S_{N}]\geq \operatorname*{\mathbb{E}}[S_{N}]]\leq\exp\left\{-\frac{2\operatorname*{ \mathbb{E}}[S_{N}]^{2}}{N(b-a)^{2}}\right\}.$ $$

In our case, let $X_{i}=\bm{1}\{\text{sample }i\text{ is routed to expert }j\}$, for some fixed $j\in[K]$. Note that $\operatorname*{\mathbb{P}}[X_{i}=1]=1/K$ due to Lemma C.6 and as such $E[S_{N}]=N/K$. Further, $X_{i}\in[0,1]$. Hence we obtain:

$$ $\displaystyle\operatorname*{\mathbb{P}}\left[S_{N}\geq\frac{2N}{K}\right]$ $\displaystyle\leq\exp\left\{-\frac{2(N/K)^{2}}{N}\right\}$ $\displaystyle\leq\exp\left\{-\frac{2N}{K^{2}}\right\}.$ $$

In order to upper-bound this probability with $\delta$, we obtain:

$$ $\displaystyle\delta\geq\exp\left\{-\frac{2N}{K^{2}}\right\}$ $\displaystyle\implies\ln\left(\frac{1}{\delta}\right)\leq\frac{2N}{K^{2}}$ $\displaystyle\implies\frac{K^{2}\ln\left(1/\delta\right)}{2}\leq N.$ $$

∎

###### Lemma C.8 .

Fix $\delta\in(0,1)$, $K\in\mathbb{N}$ and the routing function as described in Lemma C.6. Given $n\geq\frac{K^{2}\ln\left(1/\delta\right)}{2}$ samples with embedding dimension $m$. For every expert $i$, with probability at least $1-\delta$, there exists an MLP of width $\tilde{O}(n/mK)$ that correctly classifies the examples routed to this expert.

###### Proof.

We begin with a result from Daniely (). Consider $h$, a depth-2 network with $q$ neurons, an activation function $\phi:\mathbb{R}\rightarrow\mathbb{R}$, which is $O(1)$-Lipschitz, piecewise twice-differentiable, and satisfies $\operatorname*{\mathbb{E}}_{X\sim\mathcal{N}(0,1)}[\phi^{\prime}(X)]=0$, weights ${\mathbf{a}}_{i}\in\{\pm 1\}$ such that $\sum_{i=1}^{q}a_{i}=O(\sqrt{q})$.
Such a network, defined by $W$ and $a$, computes the function:

$$ $h_{W,a}(x)=\frac{1}{\sqrt{q}}\sum_{i=1}^{q}a_{i}\phi(\langle w_{i},x\rangle).$ $$

For our proof we consider $\phi$ to be the absolute value function, and later show how to obtain an MLP with ReLU activations.

Consider the set of samples $((x_{1},y_{1}),\ldots,(x_{N},y_{N}))$, where $x_{i}\sim\mathcal{N}(0,I_{M})$ and $y_{i}\in\{\pm 1\}$ is a Rademacher random variable. The objective is to find a $W^{*},a^{*}$ such that $y_{i}h_{W^{*},a^{*}}(x_{i})>0\ \forall i$. Daniely () gives the following result: for $N\leq\frac{Mq}{\log^{4}(M)}$ and $q\geq\log^{4}(M)$, with probability $1-o(1)$ there exists $h$ such that, for every $i\in[N]$, $y_{i}h(x)=\Omega(\ln(M))$.

We build on this result as follows: first, we assume that the attention matrix performs averaging of the input sequence, such that $x_{i}=\frac{1}{N}\sum_{j=1}^{N}(X^{i})_{j}$ where where $X^{i}$ is the $i$-th input sequence to memorize. Assuming the inputs are also Gaussian, then so is the average. Hence, $x_{i}\in\mathbb{R}^{m}$ and $x_{i}\sim\mathcal{N}(0,cI_{m})$ for $c=\frac{1}{L}$.

Next, we fix the number of experts to be $K$. We fix a router function $r:\mathbb{R}^{\log K}\rightarrow[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 C.6 where the router only acts on the first $\log K$ entries of each vector.
Hence, we have $K$ MLPs, each tasked with memorizing at most $\frac{2n}{K}$ input-output values. We use the network construction from Daniely () to only the last $m-\log K\geq m/2$ coordinates, which are not used by the router, so that their distribution is independent and Gaussian.

Hence, with $n\leq\frac{K}{2}\frac{(m-\log K)q}{\log^{4}(m-\log K))}\leq\frac{K}{4}\frac{mq
}{\log^{4}(m-\log K))}$ and $q\geq\log^{4}(m-\log K)$, with high probability, $y_{i}h_{j}(x_{i})=\Omega(\ln(m))$, where $h_{j}$ is the expert to which $x_{i}$ is routed. It follows that width

$$ $q\geq\frac{4n\log^{4}(m-\log K)}{mK}$ $$

is sufficient for each individual expert, and we assume that $\frac{4n}{mK}\geq 1$.

In order to obtain a MLP with a ReLU activation defined as $\sigma(x)=\max\{0,x\}$, we use the fact that $\lvert x\rvert=\sigma(x)+\sigma(-x)$. This is because absolute value is a valid activation function in Daniely (), but ReLU is not. However, doubling the width of each MLP is sufficient to obtain MLPs with ReLU activations instead, i.e.:

$$ $q\geq\frac{8n\log^{4}(m-\log K)}{mK}.$ $$

This is done by taking a solution which is of form:

$$ $h_{W,a}(x)=\frac{1}{\sqrt{q}}\sum_{i=1}^{q}a_{i}\lvert\langle w_{i},x\rangle\rvert,$ $$

and converting it to the form:

$$ $h_{W,a}^{{}^{\prime}}(x)=\frac{1}{\sqrt{q}}\sum_{i=1}^{q}a_{i}\left((\sigma{ \langle w_{i},x\rangle})+(\sigma{\langle-w_{i},x\rangle})\right).$ $$

The router has $K$ vectors each of dimension $m$. Consequently, we need $O(Km)$ parameters to store the router. Each expert has width $\frac{8n\log^{4}(m-\log K)}{mK}=\tilde{O}\left(\frac{n}{mK}\right)$ and $\tilde{O}\left(\frac{n}{K}\right)$ parameters.
∎

###### Corollary C.9 .

Let $\delta\in(0,1)$, and fix $K>1$. For $n\geq\frac{K^{2}\ln(K/\delta)}{2}$, with probability at least $1-\delta$, there exists a sparse transformer $s$ with $K$ experts such that $y_{i}s(x_{i})=\Omega(\ln m)$. It has $\tilde{O}(n+Km)$ parameters and $\tilde{O}(n/K+Km)$ active parameters.

###### Proof.

For each expert, we apply the result from Lemma C.8 with $\delta^{\prime}=\delta/K$. Hence, for every expert, with probability at most $\delta/K$, there does not exist an MLP of width $\tilde{O}(n/mK)$ which memorizes its samples.
By a union bound, this implies that with probability at most $\delta$, at least one of the experts is not able to memorize its samples. Hence, with probability at least $1-\delta$, all of the experts are able to perform memorization on their given samples. In total, this implies we will use $\tilde{O}(n+K)$ parameters to store the entire mixture of expert network. The number of active parameters is $\tilde{O}(n/K+K)$.
∎

###### Lemma C.10 (Bound on $\ell_{2}$ -norm of vector from $\mathcal{N}(0,cI_{m})$ ) .

Let $x\sim\mathcal{N}(0,cI_{m})$ for some $c>0$, and let $\delta\in(0,1)$. Then, with probability at least $1-\delta$, there exists a constant $c_{\delta,m}>0$, such that

$$ $\|x\|_{2}\leq c_{\delta,m}=\sqrt{c\left(m+2\sqrt{m\ln\left(\dfrac{1}{\delta} \right)}+2\ln\left(\dfrac{1}{\delta}\right)\right)}.$ $$

###### Proof.

Each component $x_{i}$ of $x$ is distributed as $x_{i}\sim\mathcal{N}(0,c)$. We can express $x_{i}$ as $x_{i}=\sqrt{c}\,z_{i}$, where $z_{i}\sim\mathcal{N}(0,1)$. Then,

$$ $\|x\|_{2}^{2}=\sum_{i=1}^{m}x_{i}^{2}=c\sum_{i=1}^{m}z_{i}^{2}.$ $$

The sum $\sum_{i=1}^{m}z_{i}^{2}$ follows a chi-squared distribution with $m$ degrees of freedom. By the Laurent-Massart theorem Laurent & Massart (), for all $t>0$, we have

$$ $\operatorname*{\mathbb{P}}\left(\sum_{i=1}^{m}z_{i}^{2}\geq m+2\sqrt{mt}+2t \right)\leq e^{-t}.$ $$

Multiplying both sides inside the probability by $c$, we obtain

$$ $\operatorname*{\mathbb{P}}\left(\|x\|_{2}^{2}\geq c\left(m+2\sqrt{mt}+2t\right )\right)\leq e^{-t}.$ $$

Setting $t=\ln\left(\dfrac{1}{\delta}\right)$, it follows that

$$ $\operatorname*{\mathbb{P}}\left(\|x\|_{2}^{2}\geq c\left(m+2\sqrt{m\ln\left( \dfrac{1}{\delta}\right)}+2\ln\left(\dfrac{1}{\delta}\right)\right)\right)\leq\delta.$ $$

Therefore, with probability at least $1-\delta$, we have

$$ $\|x\|_{2}\leq\sqrt{c\left(m+2\sqrt{m\ln\left(\dfrac{1}{\delta}\right)}+2\ln \left(\dfrac{1}{\delta}\right)\right)}.$ $$

∎

###### Lemma C.11 (Bounded bit complexity required) .

With high probability, $\tilde{O}(1)$ bits are required to store each weight in the network, and $\tilde{O}(n+K)$ bits are required to store the entire network. $\tilde{O}\left(\frac{n}{K}+K\right)$ active bits are required.

###### Proof.

We show that the learning result from Daniely () gives a network with bounded bit complexity. Hence, it suffices to bound the bit complexity of the initial weights and the bit complexity of the gradient update.
First, we begin with $h(x)=\sqrt{\frac{1}{q}}\sum_{i=1}^{q}a_{i}\sigma(\langle w_{i},x\rangle)$. The objective is to show there exists $\tilde{h}(x)$ with bounded bit complexity which satisfies $y_{i}\tilde{h}(x_{i})>0\ \forall i$. Suppose we begin by creating bins for the weights, and replace each $w_{i}$ with $\tilde{w}_{i}$ where $\|w_{i}-\tilde{w}_{i}\|_{\infty}\leq\varepsilon$. Further, assume that $\|w_{i}\|_{\infty}\leq M$ for some $M>0$. We then obtain that $\|w_{i}-\tilde{w}_{i}\|_{2}\leq\sqrt{m}\varepsilon$.
In addition, we state that with probability at least $1-\delta$, $\|x\|_{2}\leq c_{\delta,m}$ due to Lemma C.10.
Using this, we have that:

$$ $\displaystyle\lvert\sigma(\langle w_{i},x\rangle)-\sigma(\langle\tilde{w}_{i}, x\rangle)\rvert$ $\displaystyle\leq\lvert\langle w_{i}-\tilde{w}_{i},x\rangle\rvert$ $\displaystyle(\sigma\text{ is }1\text{-Lipschitz})$ $\displaystyle\leq\|w_{i}-\tilde{w}_{i}\|_{2}\|x\|_{2}$ $\displaystyle(\text{Cauchy-Schwarz})$ $\displaystyle\leq\sqrt{m}\varepsilon c_{\delta,m}\text{ w.p.}\geq 1-\delta$ $\displaystyle(\text{Lemma \ref{lemma:l2bound}}).$ $$

We apply this result to obtain that with probability at least $1-\delta$,

$$ $\displaystyle\lvert h(x)-\tilde{h}(x)\rvert$ $\displaystyle\leq\frac{1}{\sqrt{q}}\sum_{i=1}^{q}\lvert a_{i}\rvert\left\lvert \sigma(\langle w_{i},x\rangle)-\sigma(\langle\tilde{w}_{i},x\rangle)\right\rvert$ $\displaystyle\leq\sqrt{q}\sqrt{m}\varepsilon c_{\delta,m}.$ $$

Given that $h(x)\geq O(1)$, then for some arbitrarily small constant (we use $\frac{1}{4}$) we require: $\frac{1}{4}\leq\sqrt{q}\sqrt{m}\varepsilon c_{\delta,m}$, or equivalently, $\varepsilon\leq\frac{1}{4c_{\delta,m}\sqrt{qm}}$. Because $h(x)\geq O(1)$, we use this separation of $\frac{1}{4}$ to show that $\tilde{h}(x)$ remains the same sign as $h(x)$ (however this constant can be replaced with an arbitrarily small constant so as to satisfy the requirement).

Consider $w_{i}^{(0)}$ to be the initialization of $w_{i}$ prior to the gradient step. Then, because, $M=\|w_{i}\|_{\infty}$, we have that $M=\max_{i}\left\{{w_{i}^{(0)}}-\eta\frac{\partial h(x)}{\partial w_{i}}\right
\}=\max_{i}\left\{{w_{i}^{(0)}}-\frac{n\log m}{m}\frac{\partial h(x)}{\partial
w
_{i}}\right\}$. Assuming that $w_{i}^{(0)}$ is initialized with bounded $\ell_{\infty}$ norm of $1$, then we obtain that $M\leq 1+\frac{n\log m}{m\sqrt{q}}c_{\delta,m}$.
This is due to Lemma C.10 and that

$$ $\displaystyle\frac{\partial h(x)}{\partial w_{i}}$ $\displaystyle=\frac{1}{\sqrt{q}}a_{i}\sigma^{\prime}(\langle w_{i},x_{i} \rangle)x_{i}$ $\displaystyle=\begin{cases}0&\text{ if }\langle w_{i},x_{i}\rangle\leq 0\\ \frac{1}{\sqrt{q}}a_{i}x_{i}&\text{ otherwise}.\end{cases}$ $$

Using this, with high probability, we require $\log\left(\frac{2M}{\varepsilon}\right)=\log\left(8Mc_{\delta}\sqrt{qm}\right)$ bits (by replacing $\varepsilon$). Hence, with high probability, we require the number of bits to be at most:

$$ $O\left(\log\left(\frac{n\log m}{\sqrt{m}}c_{\delta,m}^{2}\right)\right)=\tilde {O}(1).$ $$

∎

We restate Theorem 3.6.

###### Theorem C.12 (Lower bound for dense model) .

Given the same task as above, a dense Transformer requires $\tilde{\Omega}(n)$ parameters to solve the memorization task.

###### Proof of Theorem 3.6 .

Let $c$ be the number of bits used for encoding each parameters (and we assume that $c$ is logarithmic in the problem parameters).
Denote by $\mathcal{H}$ the class of all transformers with $W$ parameters and $c$ bits per parameters. Since $\mathcal{H}$ is a finite class, where each function in the class can be encoded with $cW$ bits, we have $\lvert\mathcal{H}\rvert\leq 2^{cW}$.
Let $X^{1},\dots,X^{N}\in\mathbb{R}^{n\times d}$ be the $N$ input points. Assume a $\mathcal{H}$ can solve the memorization task. Then, for every choice of $y_{1},\dots,y_{N}\in\{\pm 1\}$, there exists a transformer $f\in\mathcal{H}$ s.t. $f(X_{i})=y_{i}$ for all $i\in[N]$. There are $2^{N}$ possible assignments for $y_{1},\dots y_{N}$ and therefore there are at least $2^{N}$ different functions in $\mathcal{H}$. So, we get $2^{N}\leq\lvert\mathcal{H}\rvert\leq 2^{cW}$ and therefore $W\geq N/c$.
∎

###### Lemma C.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 a sparse transformer than a dense transformer.

###### Proof.

As shown, a dense transformer requires $\tilde{\Omega}(n)$ parameters (and active parameters) to perform memorization of $N$ sequences. In contrast, for $n$ sufficiently large and fixed $K>1$, it holds that $\frac{n}{K}+K<n$, which shows that the number of active parameters required is less for a sparse transformer.
∎

###### Lemma C.6 .

###### Proof.

###### Lemma C.7 .

###### Proof.

###### Lemma C.8 .

###### Proof.

###### Corollary C.9 .

###### Proof.

###### Lemma C.10 (Bound on ℓ 2 subscript ℓ 2 \ell_{2} roman_ℓ start_POSTSUBSCRIPT 2 end_POSTSUBSCRIPT -norm of vector from 𝒩 ⁢ ( 0 , c ⁢ I m ) 𝒩 0 𝑐 subscript 𝐼 𝑚 \mathcal{N}(0,cI_{m}) caligraphic_N ( 0 , italic_c italic_I start_POSTSUBSCRIPT italic_m end_POSTSUBSCRIPT ) ) .

###### Proof.

###### Lemma C.11 (Bounded bit complexity required) .

###### Proof.

###### Theorem C.12 (Lower bound for dense model) .

###### Proof of Theorem 3.6 .

###### Lemma C.13 (Active parameter comparison between dense and sparse transformers) .

###### Proof.

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
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?

Role 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.
soundness: 3
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.

Empirical 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.

Provides 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.
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.
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.

Limited 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.

Diversify 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.

Ablation 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.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
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.
2. 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
3. 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. 
3.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.
4. 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.
5. 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)
6. 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?
7. 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.
8. 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.
soundness: 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
2. Clarity: The paper is clear and well-organized
3. Glad generative tasks like GSM8K and MATH were included.
4. Solid motivation of MoEs in the literature in intro.
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.
weaknesses: Scientific Weakness
1. 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.
2. 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).
3. 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. 
4. 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.
5. Question 1 in question section.
6. Intuitive explanation is your results, see my question 8.
7. 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.

Willing to raise my score if authors address the weakness well, ideally with experiments. It’s already sorted in order of importance/urgency. 

Writing Weakness
1. 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.
2. 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.
3. 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.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 4
questions: Why did you choose different architectures for the synthetic and non-synthetic tasks?

Did 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.
soundness: 4
strengths: Tradeoffs between dense and sparse models is a very relevant problem as it’s a differentiator between many existing foundation models.

The 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). 

The results are very convincing, as this trend is demonstrated on theoretical tasks, synthetic tasks designed for reasoning/memorization and on popular benchmarks
- Experiments in each section seem reasonable to back up the claim!

Graphs are very clear and the overall story is very easy to follow.
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.
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:
- Wouldn’t this bias show by just having lower perplexity? Why does the perplexity-accuracy relationship show this?

Many 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.

---

confidence: 4
contribution: 4
initial_rating: 8
presentation: 4
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.
soundness: 4
strengths: 1. The paper is well-written and easy to follow, the problem this paper focuses is important to the community.
2. Theoretical analysis is sound.
3. 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.
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.
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. 
2. 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?
3. 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?
4. Compared to dense transformers, if there are configurations that can perform better than dense transformers theoretically?
5. I'm curious about if the MoE performance on memory tasks and reasoning tasks is robust to the initialization or distribution of parameters.

Meta-review: (a) Summary of scientific Claims and Findings:
The paper investigates performance tradeoffs between Mixture-of-Experts (MoE) and dense transformer architectures. The key findings are:
- As the number of experts increases (while fixing active parameters), memorization/recall performance improves while reasoning capabilities saturate
- The authors prove theoretical limitations showing certain graph problems cannot be solved by MoEs of a given width, while being solvable by slightly wider dense models
- MoEs can effectively leverage a small number of active parameters with many experts for memorization/recall tasks

(b) Strengths:
1. Comprehensive analysis combining theoretical foundations, synthetic experiments, and real-world evaluations
2. Clear demonstration of architecture-specific tradeoffs with practical implications for model selection
3. Extensive experimental validation across multiple task types and benchmarks
4. Well-structured presentation with clear graphs and coherent narrative
5. Important and timely contribution given the growing use of MoE architectures

(c) Weaknesses:
1. Scale limitations - experiments only conducted up to 2.1B parameters, leaving uncertainty about generalization to larger scales
2. Limited exploration of MoE variants and routing strategies beyond standard top-2 token choice
3. Initial framing of results in terms of "memorization" rather than more precise concepts like generalization gaps and recall performance
4. Some synthetic reasoning tasks (like random graphs) may not fully capture real-world reasoning requirements
5. Could benefit from more analysis of tasks combining both recall and reasoning aspects

(d) Reasons for Accept:
1. Strong theoretical foundation complemented by comprehensive empirical validation
2. Clear and important insights about architectural tradeoffs that can guide practical model design decisions
3. Well-executed study with thorough experimental design and clear presentation
4. Authors demonstrated willingness to clarify and refine terminology (e.g., shifting from "memorization" to "recall")
5. Additional MMLU experiments during rebuttal further validated main claims

Decision: Accept (Poster)