# MoEUT: Mixture-of-Experts Universal Transformers

Róbert Csordás1,2† Kazuki Irie3 Jürgen Schmidhuber2,4
Christopher Potts1 Christopher D. Manning1
1Stanford University, Stanford, CA, USA
2The Swiss AI Lab IDSIA, USI & SUPSI, Lugano, Switzerland
3Center for Brain Science, Harvard University, Cambridge, MA, USA
4AI Initiative, KAUST, Thuwal, Saudi Arabia
<EMAIL_ADDRESS>
<EMAIL_ADDRESS><EMAIL_ADDRESS>
###### Abstract

22footnotetext: Work started at IDSIA.

Previous work on Universal Transformers (UTs) has demonstrated the importance
of parameter sharing across layers. By allowing recurrence in depth, UTs have
advantages over standard Transformers in learning compositional
generalizations, but layer-sharing comes with a practical limitation of
parameter-compute ratio: it drastically reduces the parameter count compared
to the non-shared model with the same dimensionality. Naively scaling up the
layer size to compensate for the loss of parameters makes its computational
resource requirements prohibitive. In practice, no previous work has succeeded
in proposing a shared-layer Transformer design that is competitive in
parameter count-dominated tasks such as language modeling. Here we propose
MoEUT (pronounced “moot”), an effective mixture-of-experts (MoE)-based shared-
layer Transformer architecture, which combines several recent advances in MoEs
for both feedforward and attention layers of standard Transformers together
with novel layer-normalization and grouping schemes that are specific and
crucial to UTs. The resulting UT model, for the first time, slightly
outperforms standard Transformers on language modeling tasks such as BLiMP and
PIQA, while using significantly less compute and memory.111Our code is public:
https://github.com/robertcsordas/moeut

##  1 Introduction

Transformers [1, 2] are ubiquitous neural architectures in modern machine
learning. They power large language models [3, 4, 5, 6, 7], modern image
processors [8], offline reinforcement learning agents [9], and many others.
Despite these successes, we should ask whether more optimal architectures
exist.

One important candidate is the Universal Transformer (UT, [10]). The core
characteristic of UTs is _recurrence in depth_ via _sharing parameters across
layers_. This reintroduces the expressive power of recurrence provided by
recurrent neural networks (RNNs, [11, 12, 13]). Layer sharing allows UTs to
outperform regular Transformers on compositional problems such as logical
inference tasks, while also yielding improvements on small-scale language
modeling and translation tasks. In particular, UTs have been shown to have
better compositional generalization properties [14, 15] by being able to
decompose structured problems without supervision and generalize to longer
sequences [16].222Dehghani et al. [10] also augment UTs with an additional
adaptive computation time (ACT, [17, 18]) mechanism. However, the benefits of
UTs we discuss here are purely due to layer-sharing, which, in consequence, is
the focus of this work. Our models could also optionally be augmented with ACT
but this is out of scope here. These empirical findings confirm that UTs are
more general architectures with superior generalization properties compared to
the standard Transformers, in principle. However, UTs suffer from a
fundamental problem of parameter-compute ratio: sharing the parameters among
$L$ layers of an $L$-layer Transformer—while keeping the same model
dimensionalities—results in a model with $L$ times less parameters (ignoring
the input/output layers to simply the discussion). Upscaling the size of the
layer to compensate for the loss of parameters (essentially by making it $L$
times wider) usually yields a very big layer whose computational requirements
in terms of compute and memory are prohibitive in practice [19, 20]. In sum,
despite their potential, UTs are much less compute-efficient than standard
Transformers, and thus, they are not popular for parameter-dominated tasks
such as modern language modeling. Indeed, we are not aware of any previous
work that has succeeded in developing compute-efficient UT models that yield
competitive performance compared to standard Transformers on such tasks.

Here we bring new perspectives and a solution to UTs’ fundamental compute-
parameter ratio problem. We present Mixture-of-Experts Universal Transformers
(MoEUTs, pronounced “moot”), a mixture-of-experts (MoE) architecture [21, 22,
23] for UTs enabling them to scale in a computationally and memory efficient
way. We leverage various recent advances in MoEs for both feedforward and
self-attention layers (Sec. 2.1 and 2.2), and combine them with two new
innovations: (1) _layer grouping_ , in which we recurrently stack groups of
MoE-based layers, and (2) a _peri-layernorm_ scheme (which is “in-between” the
standard pre- and post-layernorm), in which we apply layer norm only before
linear layers that immediately precede sigmoid or softmax activations. Both
are specifically designed for shared-layer MoE architectures, and strongly
supported by empirical evidence.

MoEUTs allow us to build parameter- and resource-efficient UT language models
outperforming standard Transformers with less compute and memory requirements
on all scales on which we can afford to test (up to 1B parameters). We
demonstrate its capabilities on the C4, SlimPajama, and peS2o language
modeling datasets, as well as on The Stack code generation. Our experiments
show that recurrence is essential for our models to achieve competitive
performance. We also demonstrate good zero-shot performance on downstream
tasks like BLiMP and Children’s Book Test, Lambada, HellaSwag, PIQA and ARC-E.
We will provide a link to our public code repository upon acceptance.

##  2 The MoEUT Architecture

Our MoEUT architecture is a Transformer architecture with shared layer
parameters, in which we address the parameter-compute ratio problem by using
mixture-of-experts. While there are many recent works on MoE methods for
Transformer language models (e.g., [24, 25, 26, 27, 28]), making them
competitive against their dense counterparts in parameter-equal comparisons is
known to be challenging [28]. Here we leverage recent advances in MoE methods
for both the feedforward network block (FFN, or simply MLP layer or
feedforward layer; Sec. 2.1) and the self-attention layer (Sec. 2.2) together
with two novel methods that take into account the specific properties of
shared-layer models, namely: layer grouping (Sec. 2.3) and signal propagation
(Sec. 2.4), which, taken together, are crucial for achieving effective shared-
layer MoE Transformers.

###  2.1 MoE Feedforward Blocks

To parameterize the feedforward blocks of our shared-layer Transformers by an
MoE, we use $\sigma$-MoE [28] with a few modifications. $\sigma$-MoE divides
the feedforward block into $N_{E}$ slices, called experts. Each expert has two
sets of weights, ${\bm{W}}_{1}^{e}\in\mathbb{R}^{d_{\text{model}}\times
d_{\text{expert}}}$ and
${\bm{W}}_{2}^{e}\in\mathbb{R}^{d_{\text{expert}}\times d_{\text{model}}}$,
where $e\in\\{1,\ldots,N_{E}\\}$ is the index of the expert. At each token
position $t$ , given layer input
${\bm{x}}_{t}\in\mathbb{R}^{d_{\text{model}}}$, the MoE feedforward layer
computes a score for each expert, yielding a vector
${\bm{s}}\in\mathbb{R}^{N_{E}}$ computed as:

$\displaystyle{\bm{s}}_{t}$ $\displaystyle=\sigma({\bm{x}}_{t}{\bm{W}}_{S})$
(1)

where ${\bm{W}}_{S}\in\mathbb{R}^{d_{\text{model}}\times N_{E}}$ is a
trainable weight matrix, and $\sigma(x)=\frac{1}{1+e^{-x}}$ is the element-
wise sigmoid function. The MoE layer only selects $K$ experts (out of $N_{E}$)
corresponding to the top-$K$ elements in ${\bm{s}}_{t}\in\mathbb{R}^{N_{E}}$
to produce the layer output ${\bm{y}}_{t}\in\mathbb{R}^{d_{\text{model}}}$ as
follows:

$\displaystyle\mathcal{E}({{\bm{x}}_{t}})$
$\displaystyle=\operatorname*{arg\,topk}({\bm{s}}_{t},K)\subseteq\\{1,\ldots,N_{E}\\}$
(2) $\displaystyle{\bm{y}}_{t}$
$\displaystyle=\sum_{e\in\mathcal{E}({\bm{x}}_{t})}{\bm{s}}_{t}[e]\operatorname{ReLU}({\bm{x}}_{t}{\bm{W}}_{1}^{e}){\bm{W}}_{2}^{e}$
(3)

where ${\bm{s}}_{t}[e]\in\mathbb{R}$ is the $e$-th element of vector
${\bm{s}}_{t}\in\mathbb{R}^{N_{E}}$. Our preliminary experiments revealed that
the original regularization of $\sigma$-MoE tends to be unstable and sometimes
causes loss explosion during training. To avoid this, we apply regularization
only within the sequence (as opposed to all tokens in the batch). For a
sequence of inputs ${\bm{x}}_{t}$, $t\in\\{1,\ldots,T\\}$ we compute the
balancing loss $L$ as:

$\displaystyle L$
$\displaystyle=\sum_{e=1}^{N_{E}}{\bm{p}}[e]\log{\bm{p}}[e],\hskip
14.22636pt{\bm{p}}=\frac{1}{T}\sum_{t=1}^{T}\mathrm{softmax}({\bm{x}}_{t}{\bm{W}}_{S})\in\mathbb{R}^{N_{E}}$
(4)

The loss is scaled with coefficient $\gamma$ and added to the standard cross
entropy loss. Unlike the original $\sigma$-MoE, no expert dropout is used in
our experiments. It is important to note that, in contrast to the standard
setup in the MoE literature, our experts are small ($d_{\text{expert}}=128$,
similarly to $\sigma$-MoE [28]), and there are 100s of them. This
configuration is called fine-grained mixture-of-experts [29] and is also
advocated by Dai et al. [30]. We analyze the effect of $d_{\text{expert}}$ in
Fig. 13 in the appendix.

###  2.2 MoE Self-Attention Layers

To introduce MoE to the self-attention layers, we apply SwitchHead [31], which
is an MoE method extending $\sigma$-MoE to attention layers. As in the
standard multi-head attention layer, each head in the SwitchHead layer
contains four transformations: query, key, value, and output projections.
However, SwitchHead parameterizes the value and output projections using MoEs.
That is, each head has one query and key projection associated with it and
$N_{A}$ value and output projections, which are chosen dynamically for each
input. Keys and queries are computed “as usual”: given an input at position
$t$, ${\bm{x}}_{t}\in\mathbb{R}^{d_{\text{model}}}$,
${\bm{k}}_{t}^{h}={\bm{x}}_{t}{\bm{W}}_{K}^{h}$ and
${\bm{q}}_{t}^{h}={\bm{x}}_{t}{\bm{W}}_{Q}^{h}$, where ${\bm{W}}_{K}^{h}$ and
${\bm{W}}_{Q}^{h}\in\mathbb{R}^{d_{\text{model}}\times d_{\text{head}}}$,
where $h\in\\{1,\ldots,{H}\\}$ is the head index. The expert selection for the
values is computed as follows:

$\displaystyle{\bm{s}}_{V,t}^{h}$
$\displaystyle=\sigma({\bm{x}}_{t}{\bm{W}}_{\mathit{SV}}^{h})\in\mathbb{R}^{N_{A}}$
(5) $\displaystyle\mathcal{E}_{V}^{h}({\bm{x}}_{t})$
$\displaystyle=\operatorname*{arg\,topk}({\bm{s}}_{V,t}^{h},K_{A})\subseteq\\{1,\ldots,N_{A}\\}$
(6)

where ${\bm{W}}_{\mathit{SV}}^{h}\in\mathbb{R}^{d_{\text{model}}\times N_{A}}$
is the selection weight for the value and $K_{A}$ is the number of
simultaneously active experts per head, set to $K_{A}=2$ in all of our
experiments. The selection for the values and outputs are independent. The
selection of the output is computed analogously using a different weight
matrix ${\bm{W}}_{\mathit{SO}}^{h}\in\mathbb{R}^{d_{\text{model}}\times
N_{A}}$:
${\bm{s}}_{O,t}^{h}=\sigma({\bm{x}}_{t}{\bm{W}}_{\mathit{SO}}^{h})\in\mathbb{R}^{N_{A}}$
and
$\mathcal{E}_{O}^{h}({\bm{x}}_{t})=\operatorname*{arg\,topk}({\bm{s}}_{O,t}^{h},K_{A})\subset\\{1,\ldots,N_{A}\\}$.
Then the output ${\bm{y}}\in\mathbb{R}^{d_{\text{model}}}$ is calculated as
follows:

$\displaystyle{\bm{v}}_{t}^{h}$
$\displaystyle=\sum_{e\in\mathcal{E}_{V}^{h}({\bm{x}}_{t})}{\bm{s}}_{V,t}^{h}[e]{\bm{x}}_{t}{\bm{W}}_{V}^{h,e}\in\mathbb{R}^{d_{\text{head}}}$
(7) $\displaystyle{\bm{a}}_{t}^{h}$
$\displaystyle=\mathrm{Attention}({\bm{q}}_{t}^{h},{\bm{K}}_{t}^{h},{\bm{V}}_{t}^{h})\in\mathbb{R}^{T}$
(8) $\displaystyle{\bm{y}}_{t}$
$\displaystyle=\sum_{h=1}^{H}\sum_{e\in\mathcal{E}_{O}^{h}({\bm{x}}_{t})}{\bm{s}}_{O,t}^{h}[e]{\bm{a}}_{t}^{h}{\bm{V}}_{t}^{h}{\bm{W}}_{O}^{h,e}$
(9)

where ${\bm{W}}_{V}^{h,e}\in\mathbb{R}^{d_{\text{model}}\times
d_{\text{head}}}$ and ${\bm{W}}_{O}^{h,e}\in\mathbb{R}^{d_{\text{head}}\times
d_{\text{model}}}$ are head $h$, expert $e$, weight matrices for value and
output respectively;
${\bm{s}}_{V,t}^{h}[e],{\bm{s}}_{O,t}^{h}[e]\in\mathbb{R}$ are scores of
expert $e$ for head $h$ at position $t$ for value and output MoE respectively;
and $\mathrm{Attention}$ denotes the standard softmax scaled dot attention [1]
with
${\bm{K}}_{t}^{h}=({\bm{k}}^{h}_{1},\ldots,{\bm{k}}^{h}_{t}),{\bm{V}}_{t}^{h}=({\bm{v}}^{h}_{1},\ldots,{\bm{v}}^{h}_{t})\in\mathbb{R}^{T\times
d_{\text{head}}}$ e.g., for the auto-regressive setting. Note that, here we
describe position-wise computations for clarity but in practice, they can be
parallelized over the tokens through matrix operations. Unlike in the original
SwitcHead, which uses no regularization, we apply the same entropy
regularization we use in the feedforward layer (Eq. 4) with a regularization
coefficient $\delta$. (The same value is used for both value and output.)

Figure 1: Layer grouping: 8 layers with group size of 2.

Figure 2: The residual grows in pre-layernorm transformers.

Figure 3: MoEUT block with no layernorms in the residual.

###  2.3 Layer Grouping: MoE-efficient Layer Sharing & Sub-operations within
an Operation

Even when using the two recent MoE methods above, which have been shown to be
successful for the standard Transformer (Sec. 2.1 and 2.2), we experimentally
observe that naive MoE-based UTs with a single shared layer often struggle to
achieve good performance at larger scales. We hypothesize that the reason is
twofold. First, as the network scales, the number of experts in the layer
grows rapidly, but we cannot increase the number of active experts $K$ at the
same rate without greatly increasing the required compute. This forces us to
reduce the percentage of active experts, which is generally detrimental.
Second, the total number of attention heads is kept relatively low, which
might not be sufficient for a large model. Increasing their number is
similarly prohibitively expensive.

Our solution to these problems is to stack multiple layers with non-shared
weights to form what we call a group of layers, reducing the number of experts
in each $\sigma$-MoE while increasing the total number of attention heads. The
final network is obtained by recurrently stacking such groups that share the
same parameters (in a sense, redefining the group as a shared “layer” in the
UT). Figure 3 provides an illustration; here, all layers denoted by “Layer A”
(or “B” respectively) share the same parameters across the entire network. The
size of the group, $G$, is the number of non-shared layers in it. In our
experiments, the group size is between 2 and 4, and the typical number of
recurrent steps is 8 or 9.

As further observations in favor of the potential inductive bias introduced by
such grouping, note that in a seminal work, Olsson et al. [32] reverse
engineer one of the main mechanisms behind in-context learning: induction
heads. They find that two successive layers where the attention performs
different operations in each layer are required. Furthermore, Csordás et al.
[16] also show that their shared-layer Transformers use two consecutive layers
to perform a _single_ operation for relatively complex synthetic tasks, such
as ListOps. Both of these observations indicate that the adjacent layers in
Transformers often perform different sub-operations for a single high-level
step of computation that spans multiple layers. This is well aligned with our
proposed grouping.

###  2.4 Novel LayerNorm Scheme for Improved Signal Propagation in Universal
Transformers

Virtually all modern Transformers make use of the so-called “pre-layernorm”
scheme [33, 34] (as opposed to the “post-layernorm” one), that is, layer
normalization [35] is applied before the attention layer (or analogously, the
feedforward block), and their output is directly added to the residual. The
residual is normalized only before the final classification layer. This design
encourages better gradient flow and is often crucial for training deep models.
This indicates that the norm of the residual vector should grow as we go
deeper in the network (see Fig. 3 for an illustration). However, it is
typically assumed that the information is carried in the _direction_ of the
residual vector instead of its length [36, 37]. Because of this, late layers
must learn to produce outputs with a larger norm so that they can apply the
same order of modification to the residual as the earlier ones, despite having
normalized inputs because of the layernorm.

These learning targets are easily achieved by standard Transformers, as they
have separate parameters which can have different scalings in different
layers, and this can be observed empirically (for more details, see Appendix
A.2). This is not the case for UTs as they have a single, shared layer (or in
our case multiple, repeated layers; see Sec. 2.3). If some circuits should be
(re-)used in both early and late layers, scaling their output to compensate
for the norm growth of the residual is nontrivial.

Post-layernorm does not have this problem, since the whole residual is
normalized after each layer. This coincides with the observation of Tan et al.
[38] that post-layernorm performs better for UTs than pre-layernorm, and with
the fact that the original UT [10] is trained with post-layernorm. That said,
as mentioned above, post-layernorm also has its own limitation in terms of
gradient flow [34].

Here we propose an alternative method to avoid the aforementioned problems: we
do not use layernorms in the “main data path”. This means, for our UTs, that
we apply no layernorm before the value projection of the attention and no
layernorm before the $\sigma$-MoE layer. Rather, layernorm is used only before
linear layers that are immediately followed by a sigmoid or softmax activation
function (producing renormalized activations that are critical before these
nonlinear layers), namely: the query and key projections in the attention, the
expert selection on both the attention and feedforward layers, and before the
final classification layer. This is illustrated in Fig. 3. Since only a ReLU
activation function is used on the main data path inside the feedforward
layer, the output updates will be proportional to the input, thus effectively
solving the residual growth issue while also providing efficient gradient flow
paths. We call this the “peri-layernorm” scheme as a scheme “between” pre- and
post-layernorm, which positions layernorm “around” (but not on) the residual
connections.

##  3 Main Experimental Results

(a)

(b)

Figure 4: Scaling of different models on C4 (with perplexity measured on a
held-out subset of C4)

We present our main experimental results on the performance and efficiency of
our MoEUT on language modeling using the popular C4 dataset [39]. To
demonstrate the versatility of our model, we also show our main results on the
SlimPajama [40] and peS2o [41] language modeling datasets, and code generation
on “The Stack” [42]. For experimental evidence in support of the benefits of
shared layers for compositional generalization, we refer to much previous work
(e.g., [10, 15, 14, 16, 38]). Following prior work [27, 31], we measure the
compute requirements in terms of the number of multiply-accumulate (MAC)
operations needed in the forward pass.

Because our models are fully MoE, they decouple the number of parameters,
compute and memory requirements, and different model dimensions such as
$d_{\text{model}}$ and $d_{\text{ff}}$, number of layers. Thus, they provide
greater flexibility for model designers. We follow a simple procedure for
setting the model’s hyperparameters, as described below. All our models use
RoPE positional encodings [43] with PyTorch’s fast attention implementation.
The baseline models are pre-layernorm Transformers. For each baseline, we
construct a parameter-matched MoEUT model. We set $d_{\text{model}}$ and the
number of layers $n_{\text{layers}}$ to be the same as for the dense baseline.
We use the same tokenization for each model trained on the same dataset. The
number of heads $H$ for MoEUT is set to $\frac{1}{4}H$ of the corresponding
dense model, $d_{\text{head}}$ is set to $2d_{\text{head}}$ of the
corresponding dense model, and we set $K_{A}=2$. This matches the number of
MACs spent for the value and output projections in self-attention, and reduces
the number of MACs spent on calculating keys and queries and the attention
matrices itself. For the $\sigma$-MoE layers, we set the expert size
$d_{\text{expert}}=128$, and $K=2d_{\text{model}}/d_{\text{expert}}$. This
halves the MAC requirements compared to the dense counterpart. We set the
number of experts in the feedforward block, $N_{E}$, and the number of
attention experts, $N_{A}$ such that the number of parameters is the same as
for the dense baseline, and $10{-}15\%$ of the model’s parameter budget
(excluding the embedding and classification layers) is spent in the attention
computations. We set the group size $G$ to 2 for all our models below 300M
parameters, $G=3$ for our 319M parameter model, and $G=4$ for the bigger
models. This helps keep the number of experts manageable and improves both the
performance and the speed of the model. All models are trained with batch size
64, context length 1024, for $10^{5}$ steps. This protocol allows us to
perform fair comparisons between different models within our computational
budget, and it leads to high quality models, as measured by our benchmarks.
For more details, see Appendix A.3.

#### Scaling compared to standard Transformers.

Our main scaling results are shown in Fig. 4. The y-axis shows the perplexity
on a held-out subset of C4. The plot shows that our MoEUT model slightly
outperforms dense models with the same number of parameters (Fig. 4(a)), and
the gap tends to grow with scale. Additionally, we compare to the non-shared
$\sigma$-MoE model [28]. This $\sigma$-MoE baseline has the same shape of
feedforward layers ($d_{\text{model}}$, $K$, $d_{\text{expert}}$) to our
layer-shared MoEUT, but uses no attention experts to keep the proportion of
the attention weights as close MoEUT as possible, and it also uses our peri-
layernorm scheme (Sec. 2.4). We add this baseline as the model that is as
close to our shared-layer model as possible. This model performs significantly
worse than MoEUT, demonstrating the clear advantage of the shared layers.
Additionally, Fig. 4(b) shows that in terms of the number of total MAC
operations spent on all forward passes during training, MoEUT outperforms the
baseline dense model by a large margin.

#### Performance on code generation.

To confirm the effectiveness of our model on a different task domain, here we
train it on a subset of the “The Stack” dataset [42] which is a code
generation task. As we cannot afford a full epoch of training, we limit
ourselves to a few languages only. We use a mixture of diverse languages:
Python, HTML, C++, Rust, JavaScript, Haskell, Scala, and assembly. We evaluate
our models on a held-out subset of the dataset. The results are shown in
Figure 6, and they are in line with our findings on the natural language
domain: MoEUT outperforms the baseline.

#### Zero-shot performance on downstream tasks.

Here we evaluate the zero-shot performance of our models on six different
downstream tasks: LAMBADA [44], BLiMP [45], Children’s Book Test (CBT) [46],
HellaSwag [47], PIQA [48] and ARC-E [49]. For LAMBADA, we use the detokenized
version from OpenAI, and we evaluate the top-1 accuracy of the last word (it
can span multiple tokens; here we use greedy decoding). For CBT and BLiMP, we
measure the accuracy for each task and report the average of the tasks’
accuracies. The results are shown in Table 1. We observe that our models and
the baselines typically perform very similarly. MoEUT often outperforms the
baseline, but the differences are marginal in all cases. This confirms that
our models are indeed very capable compared to standard language models. We
confirm this on peS2o and SlimPajama datasets as well.

Figure 5: Performance of MoEUT compared to a standard Transformer on The
Stack.

Figure 6: Perplexity of 244M MoEUT model with different layer grouping
options.

Table 1: Zero-shot downstream performance and perplexity on various language modeling datasets. Dataset | #params | Model |  PPL $\downarrow$ |  LAMBADA $\uparrow$ |  BLiMP $\uparrow$ |  CBT $\uparrow$ |  HellaSwag $\uparrow$ |  PIQA $\uparrow$ |  ARC-E $\uparrow$ |  Average $\uparrow$
---|---|---|---|---|---|---|---|---|---|---
C4 | 44M | Baseline | 18.97 | 21.9% | 73.5% | 81.3% | 28.3% | 59.9% | 31.7% | 49.4%
MoEUT |  18.30 |  23.2% |  78.2% |  81.1% |  29.2% |  61.3% |  33.5% |  51.1%
126M | Baseline | 14.97 |  28.5% | 77.0% |  84.4% | 31.7% | 62.7% | 35.2% | 53.2%
MoEUT |  14.76 | 27.2% |  79.4% | 84.2% |  32.3% |  64.4% |  35.3% |  53.8%
244M | Baseline | 13.40 |  33.1% | 78.5% |  86.0% | 34.5% | 64.9% |  36.9% |  55.6%
MoEUT |  13.24 | 30.6% |  79.7% | 85.3% |  35.7% |  65.2% | 36.4% | 55.5%
319M | Baseline | 12.81 |  33.3% | 78.5% |  87.2% | 36.1% |  67.1% | 37.2% |  56.6%
MoEUT |  12.65 | 30.8% |  80.2% | 86.9% |  37.3% | 67.0% |  37.3% |  56.6%
728M | Baseline | 11.59 |  37.8% | 80.7% | 88.2% | 40.5% | 67.7% | 39.3% | 59.0%
MoEUT |  11.34 | 36.0% |  80.8% |  88.4% |  41.8% |  69.2% |  39.6% |  59.3%
1040M | Baseline | 11.15 |  38.4% | 81.2% | 89.0% | 42.0% | 68.6% | 39.7% | 59.8%
MoEUT |  10.90 |  38.4% |  81.6% |  89.2% |  43.7% |  69.9% |  41.3% |  60.7%
peS2o | 44M | Baseline | 11.46 |  13.2% | 66.5% | 68.6% |  28.5% |  56.3% |  32.0% | 44.2%
MoEUT |  11.09 | 13.1% |  68.7% |  69.6% | 28.3% | 55.1% | 31.4% |  44.4%
244M | Baseline | 8.55 | 18.7% | 72.8% |  78.0% |  30.4% |  56.3% | 35.0% | 48.5%
MoEUT |  8.52 |  19.4% |  73.5% | 77.4% | 30.1% |  56.3% |  35.6% |  48.7%
SlimPajama | 44M | Baseline | 16.42 |  20.0% | 72.8% | 80.7% | 27.5% | 57.0% | 31.6% | 48.3%
MoEUT |  15.77 | 19.8% |  75.9% |  82.1% |  28.0% |  57.5% |  32.1% |  49.2%
244M | Baseline | 11.51 |  31.9% | 78.6% |  87.3% | 31.7% | 60.9% |  36.6% |  54.5%
MoEUT |  11.47 | 30.7% |  80.2% | 86.8% |  32.0% |  61.7% | 35.8% |  54.5%
1040M | Baseline | 9.56 |  38.8% | 80.5% | 89.9% | 37.6% | 64.5% | 38.7% | 58.3%
MoEUT |  9.36 | 38.0% |  82.5% |  90.2% |  38.1% |  64.6% |  39.1% |  58.7%

#### Comparing with SUT.

Here we compare our MoEUT to another baseline, Sparse Universal Transformer
(SUT; [38]) which is a recently proposed UT model that also makes use of MoE
layers. We note that SUTs have not been evaluated previously on any language
modeling tasks. While both MoEUT and SUT make use of an MoE for both
feedforward and attention layers, there are several technical differences at
various levels between the two methods: SUT uses competitive expert selection
(softmax), multiple load balancing losses, and much bigger expert sizes. Their
model is post-layernorm and does not use layer grouping. Unlike ours, the ACT
is used in the layer dimension.

We took the original code released by Tan et al. [38] and ported it to our
training pipeline for a fair comparison. As for MoEUT, we roughly match the
model’s dimensionalities and number of active channels to our dense baselines.
We ran a hyperparameter optimization for the regularization losses, and we
found that a minimal regularization is necessary for stabilizing the training.
However, larger regularization tends to hurt performance significantly. For
any other hyperparameters, we set them based on their biggest translation
experiments. The results are shown in Fig. 8. Effectively, SUTs, which lack
our specific methods, have a significant performance disadvantage compared to
our MoEUT and the parameter matched dense baseline. To the best of our
knowledge, we are not aware of any prior UT architectures that are both
competitive and efficient in language modeling.

#### Evaluating layer grouping.

We investigate the effect of the layer grouping (Sec. 2.3) on our 244M
parameter MoEUT model in Fig. 6. Here, $G$ denotes the number of non-shared
layers within the group. $G=2$ corresponds to the model used in all other
analyses. $G=1$ is a fully shared-layer model, without any grouping, and
$G=18$ corresponds to the baseline fully _non-shared_ $\sigma$-MoE model [28].
All hyperparameters are identical among all models, except for the number of
MLP experts ($N_{E}$) and attention experts ($N_{A}$), which are adjusted to
match the parameter count of the dense baseline. In Fig. 6, we observe that
$G=2$ is optimal, and the recurrence in the layer dimension is indeed
beneficial. Another interesting question is whether the grouping described in
Sec. 2.3 and Fig. 3 is the right way to stack layers. Let us call the two
layers in the group, A and B. The grouping we discussed so far stacks layers
in the form of “ABABAB”, e.g., for a 6-layer network. An alternative is to
first repeat one of the layers multiple times, followed by the repeated
version of the other: “AAABBB”. The “AABB” column of Fig. 6 shows this setup
for our best $G=2$ model. It can be seen that the grouping proposed in Sec.
2.3 indeed works significantly better. In fact, the AABB-style stacking is
almost as bad as not doing grouping at all.

#### Evaluating layernorm schemes.

Here we evaluate our “peri-layernorm” scheme (Sec. 2.4). Fig. 8 shows the
results. The proposed layernorm scheme consistently performs the best. The gap
is more significant for the small models, while for the bigger ones the gains
diminish (for the 719M-parameter model, the gap between peri-norm and post-
norm is marginal: 11.29 perplexity points compared to 11.32). At the same
time, we also observe that the gap between peri-norm and post-norm increases
with the number of training steps, leaving open the possibility of higher
gains if the models are trained longer. All our MoEUT models in other
experiments make use of this peri-norm scheme.

Figure 7: Comparing MoEUT and SUT [38]

Figure 8: Comparing layernorm variants.

##  4 Analysis

Here we aim to better understand the learned expert selection of MoEUTs. In
what follows, we analyze the expert selection in the MLP blocks (Sec. 2.1) of
our 244M parameter MoEUT model trained on C4. All experiments in this section
are performed by calculating statistics on the validation set of C4 for a
model with $G=2$ (i.e., two layers in the group; see Sec. 2.3). We only
display behaviors of the first layer of the group, as we find that the results
of the second layer are qualitatively similar.

Expert (re)use across layers. We first focus on whether nontrivial reuse
occurs between different layers. Note that MoE-based shared-layer models
could, in theory, assign different experts to different layers to “emulate”
regular non-shared models. If this was the case, the model would resemble a
regular Transformer instead of a Universal one. To confirm that our model is
more versatile than that, we analyze whether certain experts are activated
only in specific layers. We measure how many times each expert is activated in
each layer. This allows us to visualize the distribution of layers that each
expert prefers. To better visualize the structure, we reorder experts by a
heuristic which we call “layer position score” defined as the average of the
layer indices weighted by the number of expert activations in that layer. Fig.
11 shows the results. The yellow spot in the bottom right corner indicates
that some experts are assigned mostly to the final layer. However, for the
other experts, there is a wide range of layers where the expert is activated.
Experts seem to be active in a continuous sequence of layers. This can be seen
in the wide, vertically lined structure. We can conclude that MoEUT is capable
of specializing in a specific layer if necessary and sharing weights between
them when advantageous.

Per-token expert selection diversity. Here we analyze the diversity of expert
selection for a given input token across different layers and contexts. For
this we measure the total number of unique experts activated for individual
tokens at different layers across different positions/contexts. The result is
shown in Fig. 11. On the x-axis, the tokens are ordered differently for each
layer based on the number of experts used in that layer. We display the most
“specialized” 1000 tokens. The minimum possible number of active experts is
16, corresponding to $K$. If only $K$ experts are consistently used by a token
across different contexts, it means that the token is fully “specialized” to
consistently use a single set of $K$ experts. This is almost the case for many
tokens if we look at the first layer (blue curve): the number of unique
experts used is low, i.e., the selection mainly depends on the token’s
identity. However, in the subsequent layers, the situation is quite different:
the total number of experts used increases significantly, indicating that the
context of the token is taken into account for the expert selection. The
diversity of the experts used peaks in the middle layers (here Layer 9) and
falls slightly for layers closer to the output. For the converse analysis of
expert specialization to tokens/layers, we refer to Appendix A.5.

Figure 9: Instance-level average expert selection similarity between layers.

Expert selection dynamics in individual columns/positions. So far, all the
results have been cumulative statistics in different input sequences and
positions. We might wonder about the selection behavior for the _individual_
Transformer columns. Is expert selection mostly constant throughout the layers
for individual columns of MoEUT? To answer this question, we calculate the
pairwise intersection-over-union of the set of selected experts between all
layers in _individual columns_ and average this metric over the whole
validation set. We show the result in Fig. 9. There is a non-negligible
overlap between the selected experts in subsequent layers; however, it is far
from complete overlap. This indicates that experts usually change dynamically
in a single column, performing different functionality in different layers.

Overall, our analysis suggests that MoEUT is capable of dynamically adapting
its expert selection mechanism to a diverse set of circumstances. Sometimes,
experts are assigned to popular tokens, while in other cases, they are shared
or specialized between layers, depending on what is better for the task.

Figure 10: Layer preference of different experts.

Figure 11: No. of experts used in different layers.

##  5 Discussion and Limitations

#### More background on UT.

We emphasize that our focus is on developing scalable and performant Universal
Transformers for parameter-dominating tasks such as language modeling. This
has been a long-standing limitation of UTs. For example, Lan et al. [50] study
shared-layer BERT models [51] and find that layer-sharing hurts performance in
exchange for better parameter efficiency. Kaplan et al. [19] report that even
though Transformer language models with inter-layer parameter-sharing scale
better in terms of number of parameters, they fail to achieve compute-
efficiency—in contrast, our MoE-based approach is more compute-efficient than
the corresponding dense baseline. On the other hand, UTs have been well-known
for their compositional generalization capabilities. We refer the readers to
the numerous corresponding works (e.g., [16, 15, 38, 52, 14, 10]) for results
supporting the benefits of layer sharing. Future work may also use our MoEUT
in such compositional settings as a generally more efficient UT architecture.

#### MoE for Transformer language models.

MoE methods for Transformer language models have seen many recent advances. It
is worth noting that despite many works on MoEs for Transformers (see, e.g.,
[24, 25, 26, 27, 28]), many of them have only focused on applying MoE to the
feedforward layers. Notable exceptions are mixture-of-attention [27] and
SwitchHead [31] (Sec. 2.2) which focus on MoE self-attention layers. In
addition, until recently, it has been considered challenging to make MoE-based
language models competitive against the dense baseline in the parameter-
matched setting (unlike FLOPs/MAC-matched settings). In our MoEUT, we use
$\sigma$-MoE [28], an MoE design that has been shown to be competitive even in
such a setting.

#### Further related works on LayerNorm and layer grouping.

There are other works that are closely related to ours regarding certain
aspects of our model. Regarding the signal propagation and layernorm [35] in
Transformers (Sec. 2.4), Xie et al. [53] analyze the growing residual norm in
standard Transformers, and propose a dual, hybrid residual stream as a remedy.
Regarding layer grouping, Takase and Kiyono [54] study various layer grouping
variants to improve the efficiency of shared-layer Transformers, also showing
that layer grouping outperforms vanilla Universal Transformers. However, they
consider models with large group sizes ($G=6$ for 12 layers) and few recurrent
steps (2). We find that models with smaller $G$ with more steps perform
better. Sometimes, layer grouping is also used to up-scaling pretrained models
[55].

Limitation/Implementation. Our current implementation of the MoE layers uses
the Triton kernel released with $\sigma$-MoE [31] for both the attention and
the MLP parts of the model. This implementation is known to be suboptimal
[31]. Compared to the standard Transformer with FlashAttention [56], our MoEUT
model trains 1.5–2x slower. We estimate that with a more optimal
implementation, the training speed should be close to the dense model, while
inference should run faster.

Massive scaling. Our experiments used a modest training regime. This led to
good models and allowed us to make rigorous comparisons, but scaling to
massive training regimes for MoEUT remains an important avenue for future
research. Such experiments would inevitably require a very large compute
cluster, but the costs could also be mitigated somewhat by work optimizing our
CUDA kernel.

##  6 Conclusion

We present MoEUT, a novel Mixture-of-Expert-based Universal Transformer (UT)
model that addresses the fundamental limitation of the standard UT in terms of
parameter-compute efficiency. MoEUT combines the most advanced MoE techniques
for Transformers with our novel layer grouping method and layernorm scheme
which are both shown to be crucial for shared-layer models. Our MoEUT allows
for training competitive UTs on parameter-dominated tasks such as language
modeling, while being significantly less compute intensive than the baselines
without layer sharing. We break this long standing limitation of UTs for the
first time. Experimentally our model outperforms dense baselines from 44M to
1B parameter scale on C4, SlimPajama, peS2o, and The Stack datasets. Zero-shot
experiments confirm that the performance of MoEUT holds on downstream tasks,
including BLiMP, CBT, Lambada, HellaSwag, PIQA and ARC-E. We hope that this
work could revive research interest in Universal Transformers at larger
scales, and serve as a stepping stone for achieving the superior
generalization properties of UTs (typically limited to synthetic problems for
now) in real-world settings.

## Acknowledgements

Christopher D. Manning is a CIFAR Fellow. This research was partially funded
by ERC Advanced grant no: 742870, project AlgoRNN. We are thankful for
hardware donations from NVIDIA and IBM. We are thankful to IDSIA for providing
part of the compute used for this project even after the authors left the lab.

## References

  * Vaswani et al. [2017] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  In _Proc. Advances in Neural Information Processing Systems (NIPS)_ , pages 5998–6008, Long Beach, CA, USA, December 2017.
  * Schmidhuber [1992] Jürgen Schmidhuber.  Learning to control fast-weight memories: An alternative to recurrent nets.  _Neural Computation_ , 4(1):131–139, 1992.
  * Radford et al. [2019] Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever.  Language models are unsupervised multitask learners.  2019\.
  * Brown et al. [2020] Tom B Brown et al.  Language models are few-shot learners.  In _Proc. Advances in Neural Information Processing Systems (NeurIPS)_ , Virtual only, December 2020.
  * OpenAI [2022] OpenAI.  ChatGPT.  https://openai.com/blog/chatgpt, 2022.
  * OpenAI [2023] OpenAI.  GPT-4 technical report.  _Preprint arXiv:2303.08774_ , 2023.
  * Touvron et al. [2023] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurélien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample.  LLaMA: Open and efficient foundation language models.  _Preprint arXiv:2302.13971_ , 2023.
  * Dosovitskiy et al. [2021] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby.  An image is worth 16x16 words: Transformers for image recognition at scale.  In _Proc. Advances in Neural Information Processing Systems (NeurIPS)_ , Virtual only, May 2021.
  * Chen et al. [2021] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Michael Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch.  Decision transformer: Reinforcement learning via sequence modeling.  In _Proc. Advances in Neural Information Processing Systems (NeurIPS)_ , pages 15084–15097, Virtual only, December 2021.
  * Dehghani et al. [2019] Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser.  Universal Transformers.  In _Int. Conf. on Learning Representations (ICLR)_ , New Orleans, LA, USA, May 2019.
  * Elman [1990] Jeffrey L. Elman.  Finding structure in time.  _Cognitive Science_ , 14(2):179–211, 1990.
  * Jordan [1986] Michael I Jordan.  Attractor dynamics and parallelism in a connectionist sequential machine.  In _Proc. Conf. of the Cognitive Science Society_ , pages 531–546. Amherst, MA, USA, August 1986.
  * Hochreiter and Schmidhuber [1997] Sepp Hochreiter and Jürgen Schmidhuber.  Long short-term memory.  _Neural computation_ , pages 1735–1780, 1997.
  * Ontañón et al. [2022] Santiago Ontañón, Joshua Ainslie, Zachary Fisher, and Vaclav Cvicek.  Making transformers solve compositional tasks.  In _Proc. Association for Computational Linguistics (ACL)_ , pages 3591–3607, Dublin, Ireland, May 2022.
  * Csordás et al. [2021] Róbert Csordás, Kazuki Irie, and Jürgen Schmidhuber.  The devil is in the detail: Simple tricks improve systematic generalization of Transformers.  In _Proc. Conf. on Empirical Methods in Natural Language Processing (EMNLP)_ , Punta Cana, Dominican Republic, November 2021.
  * Csordás et al. [2022] Róbert Csordás, Kazuki Irie, and Jürgen Schmidhuber.  The neural data router: Adaptive control flow in transformers improves systematic generalization.  In _Int. Conf. on Learning Representations (ICLR)_ , Virtual only, April 2022.
  * Schmidhuber [2012] Jürgen Schmidhuber.  Self-delimiting neural networks.  _Preprint arXiv:1210.0118_ , 2012.
  * Graves [2016] Alex Graves.  Adaptive computation time for recurrent neural networks.  In _Int. Conf. on Learning Representations (ICLR) Workshop Track_ , Vancouver, Canada, April 2016.
  * Kaplan et al. [2020] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei.  Scaling laws for neural language models.  _Preprint arXiv:2001.08361_ , 2020.
  * Tay et al. [2023] Yi Tay, Mostafa Dehghani, Samira Abnar, Hyung Won Chung, William Fedus, Jinfeng Rao, Sharan Narang, Vinh Q. Tran, Dani Yogatama, and Donald Metzler.  Scaling laws vs model architectures: How does inductive bias influence scaling?  In _Findings of the Association for Computational Linguistics: EMNLP_ , Singapore, December 2023.
  * Jacobs et al. [1991] Robert A. Jacobs, Michael I. Jordan, Steven J. Nowlan, and Geoffrey E. Hinton.  Adaptive mixtures of local experts.  _Neural Compututaion_ , 3(1):79–87, 1991.
  * II and Waibel [1990] John B. Hampshire II and Alexander H. Waibel.  The meta-pi network: connectionist rapid adaptation for high-performance multi-speaker phoneme recognition.  In _Proc. IEEE Int. Conf. on Acoustics, Speech and Signal Processing (ICASSP)_ , pages 165–168, Albuquerque, New Mexico, USA, April 1990.
  * Shazeer et al. [2017] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean.  Outrageously large neural networks: The sparsely-gated mixture-of-experts layer.  In _Int. Conf. on Learning Representations (ICLR)_ , Toulon, France, April 2017.
  * Lepikhin et al. [2021] 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.  In _Int. Conf. on Learning Representations (ICLR)_ , Virtual only, May 2021.
  * Fedus et al. [2021] William Fedus, Barret Zoph, and Noam Shazeer.  Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity.  _Preprint arXiv:2101.03961_ , 2021.
  * Clark et al. [2022] Aidan Clark, Diego de Las Casas, Aurelia Guy, Arthur Mensch, Michela Paganini, Jordan Hoffmann, Bogdan Damoc, Blake A. Hechtman, Trevor Cai, Sebastian Borgeaud, George van den Driessche, Eliza Rutherford, Tom Hennigan, Matthew Johnson, Katie Millican, Albin Cassirer, Chris Jones, Elena Buchatskaya, David Budden, Laurent Sifre, Simon Osindero, Oriol Vinyals, Jack W. Rae, Erich Elsen, Koray Kavukcuoglu, and Karen Simonyan.  Unified scaling laws for routed language models.  _Preprint arXiv:2202.01169_ , 2022.
  * Zhang et al. [2022] Xiaofeng Zhang, Yikang Shen, Zeyu Huang, Jie Zhou, Wenge Rong, and Zhang Xiong.  Mixture of attention heads: Selecting attention heads per token.  In _Proc. Conf. on Empirical Methods in Natural Language Processing (EMNLP)_ , pages 4150–4162, Abu Dhabi, United Arab Emirates, December 2022.
  * Csordás et al. [2023a] Róbert Csordás, Kazuki Irie, and Jürgen Schmidhuber.  Approximating two-layer feedforward networks for efficient transformers.  In _Findings of the Association for Computational Linguistics: EMNLP 2023_ , November 2023a.
  * Krajewski et al. [2024] Jakub Krajewski, Jan Ludziejewski, Kamil Adamczewski, Maciej Pióro, Michal Krutul, Szymon Antoniak, Kamil Ciebiera, Krystian Król, Tomasz Odrzygózdz, Piotr Sankowski, Marek Cygan, and Sebastian Jaszczur.  Scaling laws for fine-grained mixture of experts.  _Preprint arXiv:2402.07871_ , 2024.
  * Dai et al. [2024] Damai Dai, Chengqi Deng, Chenggang Zhao, R. X. Xu, Huazuo Gao, Deli Chen, Jiashi Li, Wangding Zeng, Xingkai Yu, Y. Wu, Zhenda Xie, Y. K. Li, Panpan Huang, Fuli Luo, Chong Ruan, Zhifang Sui, and Wenfeng Liang.  DeepSeekMoE: Towards ultimate expert specialization in mixture-of-experts language models.  _Preprint arXiv:2401.06066_ , 2024.
  * Csordás et al. [2023b] Róbert Csordás, Piotr Piękos, Kazuki Irie, and Jürgen Schmidhuber.  SwitchHead: Accelerating transformers with mixture-of-experts attention.  In _Preprint arXiv:2312.07987_ , December 2023b.
  * Olsson et al. [2022] Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah.  In-context learning and induction heads.  _Transformer Circuits Thread_ , 2022.  https://transformer-circuits.pub/2022/in-context-learning-and-induction-heads/index.html.
  * Xiong et al. [2020] Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tie-Yan Liu.  On layer normalization in the transformer architecture.  In _Proc. Int. Conf. on Machine Learning (ICML)_ , volume 119, pages 10524–10533, Virtual Only, July 2020.
  * He et al. [2016] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Identity mappings in deep residual networks.  In _Proc. European Conf. on Computer Vision (ECCV)_ , pages 630–645, Amsterdam, Netherlands, October 2016.
  * Ba et al. [2016] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton.  Layer normalization.  _Preprint arXiv:1607.06450_ , 2016.
  * Thorpe [1989] Simon J. Thorpe.  Local vs. distributed coding.  _Intellectica_ , 8:3–40, 1989.
  * Elhage et al. [2022] Nelson Elhage, Tristan Hume, Catherine Olsson, Nicholas Schiefer, Tom Henighan, Shauna Kravec, Zac Hatfield-Dodds, Robert Lasenby, Dawn Drain, Carol Chen, Roger Grosse, Sam McCandlish, Jared Kaplan, Dario Amodei, Martin Wattenberg, and Christopher Olah.  Toy models of superposition.  _Transformer Circuits Thread_ , 2022.
  * Tan et al. [2023] Shawn Tan, Yikang Shen, Zhenfang Chen, Aaron C. Courville, and Chuang Gan.  Sparse universal transformer.  In _Proc. Conf. on Empirical Methods in Natural Language Processing (EMNLP)_ , pages 169–179, Singapore, December 2023.
  * Raffel et al. [2020] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu.  Exploring the limits of transfer learning with a unified text-to-text transformer.  _Journal of Machine Learning Research (JMLR)_ , 21:140:1–140:67, 2020.
  * Soboleva et al. [2023] Daria Soboleva, Faisal Al-Khateeb, Robert Myers, Jacob R Steeves, Joel Hestness, and Nolan Dey.  SlimPajama: A 627B token cleaned and deduplicated version of RedPajama, June 2023.  URL https://huggingface.co/datasets/cerebras/SlimPajama-627B.
  * Soldaini and Lo [2023] Luca Soldaini and Kyle Lo.  peS2o (Pretraining Efficiently on S2ORC) Dataset.  Technical report, Allen Institute for AI, 2023.  https://github.com/allenai/pes2o.
  * Kocetkov et al. [2022] Denis Kocetkov, Raymond Li, Loubna Ben Allal, Jia Li, Chenghao Mou, Carlos Muñoz Ferrandis, Yacine Jernite, Margaret Mitchell, Sean Hughes, Thomas Wolf, Dzmitry Bahdanau, Leandro von Werra, and Harm de Vries.  The Stack: 3 TB of permissively licensed source code.  _Preprint arXiv:2211.15533_ , 2022.
  * Su et al. [2021] Jianlin Su, Yu Lu, Shengfeng Pan, Bo Wen, and Yunfeng Liu.  RoFormer: Enhanced transformer with rotary position embedding.  _Preprint arXiv:2104.09864_ , 2021.
  * Paperno et al. [2016] Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Quan Ngoc Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández.  The LAMBADA dataset: Word prediction requiring a broad discourse context.  In _Proc. Association for Computational Linguistics (ACL)_ , Berlin, Germany, August 2016.
  * Warstadt et al. [2020] Alex Warstadt, Alicia Parrish, Haokun Liu, Anhad Mohananey, Wei Peng, Sheng-Fu Wang, and Samuel R. Bowman.  BLiMP: The benchmark of linguistic minimal pairs for English.  _Transactions of the Association for Computational Linguistics (TACL)_ , 8:377–392, 2020.
  * Hill et al. [2016] Felix Hill, Antoine Bordes, Sumit Chopra, and Jason Weston.  The Goldilocks principle: Reading children’s books with explicit memory representations.  In _Int. Conf. on Learning Representations (ICLR)_ , San Juan, Puerto Rico, May 2016.
  * Zellers et al. [2019] Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi.  Hellaswag: Can a machine really finish your sentence?  In _Proc. Association for Computational Linguistics (ACL)_ , pages 4791–4800, Florence, Italy, August 2019.
  * Bisk et al. [2020] Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi.  PIQA: reasoning about physical commonsense in natural language.  In _Proc. AAAI Conf. on Artificial Intelligence_ , pages 7432–7439, New York, NY, USA, February 2020. AAAI Press.
  * Clark et al. [2018] Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord.  Think you have solved question answering? try ARC, the AI2 reasoning challenge.  _Preprint arXiv:1803.05457_ , 2018.
  * Lan et al. [2020] Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut.  ALBERT: A lite BERT for self-supervised learning of language representations.  In _Int. Conf. on Learning Representations (ICLR)_ , Virtual only, April 2020.
  * Devlin et al. [2019] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  BERT: pre-training of deep bidirectional Transformers for language understanding.  In _Proc. North American Chapter of the Association for Computational Linguistics on Human Language Technologies (NAACL-HLT)_ , pages 4171–4186, Minneapolis, MN, USA, June 2019.
  * Bergen et al. [2021] Leon Bergen, Timothy J. O’Donnell, and Dzmitry Bahdanau.  Systematic generalization with edge transformers.  In _Proc. Advances in Neural Information Processing Systems (NeurIPS)_ , pages 1390–1402, Virtual only, December 2021.
  * Xie et al. [2023] Shufang Xie, Huishuai Zhang, Junliang Guo, Xu Tan, Jiang Bian, Hany Hassan Awadalla, Arul Menezes, Tao Qin, and Rui Yan.  Residual: Transformer with dual residual connections.  _Preprint arXiv:2304.14802_ , 2023.
  * Takase and Kiyono [2023] Sho Takase and Shun Kiyono.  Lessons on parameter sharing across layers in transformers.  In Nafise Sadat Moosavi, Iryna Gurevych, Yufang Hou, Gyuwan Kim, Young Jin Kim, Tal Schuster, and Ameeta Agrawal, editors, _SustaiNLP Workshop_ , pages 78–90, Toronto, Canada, July 2023. Association for Computational Linguistics.
  * Kim et al. [2023] Dahyun Kim, Chanjun Park, Sanghoon Kim, Wonsung Lee, Wonho Song, Yunsu Kim, Hyeonwoo Kim, Yungi Kim, Hyeonju Lee, Jihoo Kim, Changbae Ahn, Seonghoon Yang, Sukyung Lee, Hyunbyung Park, Gyoungjin Gim, Mikyoung Cha, Hwalsuk Lee, and Sunghun Kim.  SOLAR 10.7B: Scaling large language models with simple yet effective depth up-scaling.  _Preprint arXiv:2312.15166_ , 2023.
  * Dao et al. [2022] Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré.  FlashAttention: Fast and memory-efficient exact attention with IO-awareness.  In _Proc. Advances in Neural Information Processing Systems (NeurIPS)_ , New Orleans, Louisiana, USA, December 2022.
  * Paszke et al. [2019] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala.  PyTorch: An imperative style, high-performance deep learning library.  In _Proc. Advances in Neural Information Processing Systems (NeurIPS)_ , pages 8024–8035, Vancouver, Canada, December 2019.
  * Loshchilov and Hutter [2019] Ilya Loshchilov and Frank Hutter.  Decoupled weight decay regularization.  In _Int. Conf. on Learning Representations (ICLR)_ , New Orleans, LA, USA, May 2019.
  * Kudo and Richardson [2018] Taku Kudo and John Richardson.  Sentencepiece: A simple and language independent subword tokenizer and detokenizer for neural text processing.  In _Proc. Conf. on Empirical Methods in Natural Language Processing (EMNLP)_ , pages 66–71, Brussels, Belgium, October 2018.

##  Appendix A Appendix

###  A.1 Broader Impact Statement

We consider this work to be a foundational research paper with no direct
societal implications. However, novel research that builds on this work may
break the generalization bottleneck of current models, allowing better
reasoning. This can potentially be a jump towards Artificial General
Intelligence, which might have unforeseeable consequences, both positive and
negative. Additionally, a better implementation of our CUDA kernels might lead
to foundation models that are more efficient than current ones, which might
make them more accessible. This can be beneficial because of the reduction in
energy usage, but it might also enable easier generation of harmful content
like fake news or deepfakes.

###  A.2 Growing Residual Norm In Standard Transformers

In Sec. 2.4, we discussed the issue of the growing residual norm in the
standard Transformers. Here, we measure the $L^{2}$ norm of the difference of
the residual before and after applying a standard Transformer layer (both the
attention and the MLP block) in different layers of our 44M parameter
Transformer trained on C4. The results are visualized in Fig. 12. It can be
seen that the norm of the updates indeed grow in later layers.

Figure 12: The update magnitude of different layers in a 44M parameter
Transformer on C4.

###  A.3 Hyperparameters

All our models are trained in PyTorch [57] with a batch size of 64, context
length of 1024, for 100k iterations, a learning rate of 0.00025, AdamW
optimizer [58] with default hyperparameters, weight decay of 0.01. They are
trained on a single node in a data-parallel manner. The learning rate is
decayed to $10\%$ of its initial value using cosine decay. We use a gradient
clipping of $\kappa$ and $N_{\text{warmup}}$ linear learning rate warmup steps
(see Tab. 2). None of our models uses dropout. For the entropy regularization
of the MLP expert selection, we use $\gamma=0.01$ and for SwitchHead attention
$\delta=0.001$. Expert dropout is not used. All of our models use a
SentencePiece [59] tokenizer with 8000 tokens, trained on a subset of the
training set for the given dataset. All models are trained with mixed
precision. The hyperparameters of the SUT models can be found in Tab. 3. Note
that the meanings of the parameters are not directly analogous to ours. Please
refer to Tan et al. [38] for more details.

Table 2: Hyperparameters of different models used in our main experiments. Model | #params | $n_{\text{layers}}$ | $G$ | $d_{\text{model}}$ | $d_{\text{ff}}$ | $H$ | $N_{A}$ | $d_{\text{head}}$ | $N_{E}$ | $K$ | $N_{\text{warmup}}$ | $\kappa$
---|---|---|---|---|---|---|---|---|---|---|---|---
Baseline | 45M | 16 | - | 412 | 2053 | 10 | - | 41 | - | - | 0 | 0.1
MoEUT | 44M | 16 | 2 | 412 | - | 4 | 8 | 82 | 155 | 12 | 0 | 0.1
$\sigma$-MoE | 44M | 16 | 16 | 412 | - | 4 | 1 | 82 | 17 | 12 | 0 | 0.1
Baseline | 126M | 16 | - | 768 | 3072 | 16 | - | 48 | - | - | 4000 | 0.25
MoEUT | 126M | 18 | 2 | 768 | - | 4 | 10 | 96 | 254 | 12 | 4000 | 0.25
$\sigma$-MoE | 126M | 18 | 18 | 768 | - | 4 | 1 | 96 | 26 | 12 | 4000 | 0.25
Baseline | 244M | 18 | - | 1024 | 4110 | 16 | - | 64 | - | - | 4000 | 0.25
MoEUT | 243M | 18 | 2 | 1024 | - | 4 | 10 | 128 | 387 | 16 | 4000 | 0.25
$\sigma$-MoE | 244M | 18 | 18 | 1024 | - | 4 | 1 | 128 | 40 | 16 | 4000 | 0.25
Baseline | 319M | 24 | - | 1024 | 4110 | 16 | - | 64 | - | - | 4000 | 0.25
MoEUT | 318M | 24 | 3 | 1024 | - | 4 | 10 | 128 | 338 | 16 | 4000 | 0.25
$\sigma$-MoE | 320M | 24 | 24 | 1024 | - | 4 | 1 | 128 | 40 | 16 | 4000 | 0.25
Baseline | 729M | 36 | - | 1280 | 5120 | 20 | - | 64 | - | - | 4000 | 0.25
MoEUT | 727M | 36 | 4 | 1280 | - | 5 | 13 | 128 | 467 | 20 | 4000 | 0.25
$\sigma$-MoE | 731M | 36 | 36 | 1280 | - | 5 | 1 | 128 | 50 | 20 | 4000 | 0.25
Baseline | 1044M | 36 | - | 1536 | 6144 | 24 | - | 64 | - | - | 4000 | 0.25
MoEUT | 1040M | 36 | 4 | 1536 | - | 6 | 12 | 128 | 565 | 24 | 4000 | 0.25
Table 3: Hyperparameters of SUT models used in our experiments. #params | $n_{\text{layers}}$ | $d_{\text{model}}$ | $d_{\text{expert}}$ | $H$ | $N_{A}$ | $d_{\text{att\\_expert}}$ | $d_{\text{head}}$ | $N_{E}$ | $K$ | $\mathcal{L}_{\text{MIM}}$ | $\mathcal{L}_{\text{ACT}}$ | $N_{\text{warmup}}$ | $\kappa$
---|---|---|---|---|---|---|---|---|---|---|---|---|---
45M | 16 | 412 | 256 | 4 | 24 | 256 | 64 | 152 | 2 | 0.001 | 0.01 | 0 | 0.1
245M | 18 | 1024 | 512 | 4 | 21 | 512 | 128 | 192 | 4 | 0.01 | 0.01 | 4000 | 0.25

###  A.4 The Effects of $d_{\text{expert}}$ and $K$

Here, we analyze the effect of the expert size given a fixed amount of compute
($d_{\text{expert}}\cdot K$ being kept constant). The results are shown in
Fig. 13. It can be seen that using fine-grained mixture of experts (small
experts) is indeed critical for good performance. In our models, we use
$d_{\text{expert}}=128$. Using smaller experts significantly decreases the
compute efficiency of the Tirton kernel. Please note that these experiments
keep the number of active channels in the MLP block constant. Thus, the effect
is based purely on the dynamics of the selection mechanism.

Figure 13: Performance of our 244M MoEUT on a held-out subset of C4 with
expert sizes ($d_{\text{expert}}$) in the MLP layer.

We also analyzed the performance of our MoEUT model with different numbers of
active experts in the MLP layer. This varies the amount of compute spent in
the layer, and the number of active channels. We show the results in Fig. 14.
Increasing the number of active experts always improves performance, but the
returns are diminishing. We chose $G=16$ for our experiments because of
efficiency reasons.

Figure 14: Performance of our 244M MoEUT on a held-out subset of C4 with
different number of active experts for the $\sigma$-MoE layer.

###  A.5 Additional Analysis

Expert specialization to tokens/layers. Now conversely to the “per-token
expert selection diversity” analysis presented in the main text, we analyze
whether the experts activated by a token are layer-specific for that specific
token. First, we count the total number of experts used by each token and
measure the proportion of them that is used in each layer. The results are
shown in Fig. 15. Here, we order the tokens (x-axis) by their frequency in the
validation set (the same ordering is used for all layers). The first 6000
tokens are shown. We can see that most of the experts are present in all
layers for the most popular tokens. This can be seen from the high score in
the first few columns. We can also observe that experts tend to be more layer-
specific for the less popular tokens. In addition, the set of experts selected
in the early layers is typically less diverse than for the rest of the layers:
only a small fraction of the used experts are present there. This is
consistent with the findings for Layer 1 in Fig. 11.

Figure 15: Fraction of experts used in a specific layer out of all experts
used by a given token.

###  A.6 Compute Requirements

We report the hardware used and the required wall clock time for our main
experiments in the paper in Tab. 4. All experiments are performed in private
clusters. The duration is reported in “hh:mm” format. We report the number of
GPUs ($N_{\text{GPU}}$) used for that specific experiment (and _not_ the total
number of GPUs in the system). For the number of CPUs ($N_{\text{CPU}}$) and
RAM we report the total amount in the node, as these resources are shared
between concurrent runs.

Note that the report is generated from Weights and Biases logs. Because of
this, the duration might include effects of restarts because of SLURM
preemption and effects for occasional slowdowns of the network drive.
Additionally, in a few instances the Weights and Biases log buffer overflowed,
making it impossible to determine the number of GPUs used for the experiment.
In this case we report “??” for the experiment. Note that we only report the
resource usage of the final experiments here. We estimate that the total cost
of the failed experiments and preliminary runs is at least an order of
magnitude higher than this.

Table 4: Training hardware information for the experiments reported in the paper Model | #params | Dataset | $G$ | GPU Type | $N_{\text{GPU}}$ | $N_{\text{CPU}}$ | RAM | Duration
---|---|---|---|---|---|---|---|---
SUT | 45M | C4 | - | TITAN RTX | 4 | 8 | 125G | 29:28
SUT | 45M | peS2o | - | RTX 3090 | 5 | 24 | 188G | 28:46
SUT | 245M | C4 | - | A100-80GB | 4 | 128 | 1007G | 28:20
SUT | 245M | peS2o | - | A100-80GB | 4 | 128 | 1007G | 23:43
MoEUT PostLN | 44M | C4 | 2 | RTX 3090 | 5 | 24 | 251G | 13:42
MoEUT PostLN | 243M | C4 | 2 | V100-32GB | 8 | 40 | 503G | 23:40
$\sigma$-MoE | 44M | C4 | - | V100-16GB | ?? | 32 | 220G | 20:09
MoEUT | 44M | C4 | 2 | V100-16GB | 4 | 40 | 251G | 19:32
MoEUT | 44M | peS2o | 2 | V100-16GB | ?? | 32 | 220G | 22:04
MoEUT | 44M | SlimPajama | 2 | RTX 3090 | 5 | 24 | 188G | 15:21
MoEUT | 126M | C4 | 2 | RTX 3090 | 5 | 24 | 251G | 20:05
$\sigma$-MoE | 126M | C4 | - | RTX 3090 | 5 | 24 | 251G | 17:27
MoEUT AABB | 243M | C4 | 2 | RTX 3090 | 5 | 24 | 188G | 41:02
MoEUT | 243M | C4 | 2 | RTX 4090 | 5 | 24 | 251G | 26:03
MoEUT PreLN | 243M | C4 | 2 | V100-32GB-LS | ?? | 40 | 503G | 46:09
MoEUT | 243M | peS2o | 2 | RTX 4090 | 5 | 24 | 251G | 24:59
MoEUT | 243M | SlimPajama | 2 | RTX 3090 | 5 | 24 | 251G | 32:55
MoEUT | 243M | TheStack | 2 | RTX 4090 | 5 | 24 | 251G | 24:36
MoEUT | 244M | C4 | 1 | RTX 3090 | 5 | 24 | 251G | 36:39
MoEUT | 244M | C4 | 3 | RTX A6000 | 5 | 48 | 503G | 11:08
MoEUT | 244M | C4 | 6 | RTX 3090 | 5 | 24 | 251G | 30:06
MoEUT | 244M | C4 | 9 | RTX 4090 | 5 | 24 | 251G | 22:30
$\sigma$-MoE | 244M | C4 | - | V100-32GB-LS | 8 | 40 | 503G | 19:34
MoEUT | 318M | C4 | 3 | V100-32GB | 8 | 40 | 503G | 27:51
$\sigma$-MoE | 320M | C4 | - | RTX 3090 | 5 | 24 | 251G | 37:06
MoEUT | 727M | C4 | 4 | A100-80GB | 4 | 128 | 1007G | 53:28
MoEUT | 727M | TheStack | 4 | A100-80GB | 4 | 128 | 1007G | 38:40
$\sigma$-MoE | 731M | C4 | - | V100-32GB-LS | 8 | 40 | 503G | 52:27
MoEUT | 1040M | C4 | 4 | A100-80GB | 4 | 128 | 1007G | 74:10
MoEUT | 1040M | SlimPajama | 4 | A100-80GB | 4 | 128 | 1007G | 70:35
Transformer | 45M | C4 | - | RTX 3090 | 2 | 48 | 251G | 24:08
Transformer | 45M | peS2o | - | V100-16GB | 4 | 32 | 220G | 13:32
Transformer | 45M | SlimPajama | - | V100-16GB | 4 | 40 | 251G | 11:59
Transformer | 126M | C4 | - | RTX A6000 | 4 | 48 | 503G | 12:03
Transformer | 244M | C4 | - | V100-32GB | 8 | 40 | 503G | 14:22
Transformer | 244M | peS2o | - | RTX A6000 | 4 | 48 | 503G | 21:04
Transformer | 244M | SlimPajama | - | V100-32GB-LS | 8 | 40 | 503G | 19:08
Transformer | 244M | TheStack | - | RTX 3090 | ?? | 24 | 251G | 26:15
Transformer | 319M | C4 | - | V100-32GB | 8 | 40 | 503G | 16:36
Transformer | 729M | C4 | - | V100-32GB | 8 | 40 | 503G | 31:29
Transformer | 729M | TheStack | - | A100-80GB | 4 | 128 | 1007G | 25:39
Transformer | 1044M | C4 | - | A100-80GB | 4 | 128 | 1007G | 31:26
Transformer | 1044M | SlimPajama | - | A100-80GB | 4 | 128 | 1007G | 31:38