MoEUT: Mixture-of-Experts Universal Transformers
Authors: Róbert Csordás
Kazuki Irie
Jürgen Schmidhuber
Christopher Potts
Christopher D Manning
Venue: NeurIPS 2024 poster (2024)
Keywords: MoE
mixture of experts
universal transformers
language modeling
Primary area: deep_learning_architectures
TL;DR: We propose Mixture of Experts Universal Transformers that scales to language modeling for the first time.

Abstract: 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.

MoEUT: Mixture-of-Experts Universal Transformers
Róbert Csordás 1.2f Kazuki Irie
Jürgen Schmidhuber
Christopher Potts' Christopher D. Manning'
Stanford University, Stanford. CA, USA
The Swiss Al Lab IDSIA, USI & SUPSI, Lugano, Switzerland 3 Center for Brain Science, Harvard University, Cambridge, MA, USA
*Al Initiative, KAUST, Thuwal, Saudi Arabia (rcsordas, cgpotts, Ostannnnnntttoorrooooooooo edu kirie0fas harvard. edu, juergen@idsia.ch
Abstract
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 learning compositional generalizations, but layer-sharing comes with a practical limitation of parameter-compute ratio: it drastically reduces the parameter count compared same dimensionality. Naively scaling up the layer size to compensate for the loss of parameters makes its computational resource requirements prohibitive. In practice, previous work has succeeded in proposing shared-layer Transformer design that is competitive in parameter count-dominated tasks such as language modeling. Here we propose MoEUT (pronounced "moot"), effective mixture-of-exxprrs (MoE)-based shared-layer Transformer architecture, which combines several recent advances in MoEs for both feedforward and attention layers of standard Trans- formers 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."
non-shared model with the
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, 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. 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].2 These empirical findings confirm that UTs are more general architectures with superior
Work started at IDSIA.
'Our code is public: https: htaps://ggthune com//robertccordas/meentt
'Dehghani 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.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

generalization properties compared to standard Transformers, in principle. However, UTs suffer from a fundamental problem of parameter- an T--laye Traanformer--we while keeping the same model iimensionallties-reellle results in a model with I times fewer parameters (ignoring the input/output layers to simplify the discussion). Upscaling the size of the layer to compensate for the loss of parameters (essentially by making it I. times wider) usually yields a very big layer whose computational requirements in terms of compute and memory 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-dominater tasks such as modern language modeling. Indeed, succeeded in developing compute-effictts efficeieff UT models that yield competitive performance compared to standard Transformers on such tasks.
ratio: sharing the parameters among layers of
not aware of any previous work that has
Here we bring new perspectives and a solution to UTs' fundamental compute-parameter ratio problem. We present Mixture-of Exxperss Universal Transformers (MoEUT's, pronounced "moot"), mixture- of-experts (MoE) architecture [21, 22, 23] for UT's enabling them to scale 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) peri-layernorm scheme (which is 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 -hared-layer MoE architectures, and strongly supported by empirical evidence.
MoEUTss 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 IB parameters). We demonstrate their 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.
2 The MoEUT Architecture
Our MoEUT architecture is a Transformer architecture with shared layer parameters, in which address the parameter-compee ratio problem by using mixture-of-cxperts. 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 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 MoE [28] with a few modifications. o-MoE divides the feedforward block into Ne slices, called experts. Each expert has two sets of weights, Wf E Rd dapen and We E Rd X dmndel where E f1....,,N Ne} } is the index of the expert. At each token position t. given layer input H, E Rd the MoE feedforward layer computes a score for each expert, yielding vector s E R computed as:
$t (2.
(1)
where Ws E Rd is trainable weight matrix, and ((x) = I+6 is the element-wise sigmoid function. The MoE layer only selects К experts (out of Ng) corresponding to the top-K K elements in 8, E Ne to produce the layer output y: Ridune as follows:
E(x:) = arg topk( st, K) S fL.....NN Net
(2)
y/ Σ
se(e] (eLU(x;WW)WW
(3)
2

---

where st [e] E IR is the e-th element of vector E IR Our preliminary experiments revealed that the original regularization of o-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 Tt, E (I....,T) we compute the balancing loss I. as:
NR L 2 píc) pog
1 1 p
T softmax Ws) E E
Σ
The loss is scaled with coefficient y and added to the standard cross entropy loss. Unlike the original o-MoE, no expert dropout used in our experiments. is important to note that, in contrast to the standard setup in the MoE literature, our experts are small (deapen = 128, similarly to o-MoE [28]). and there are 100s of them. This configuration is called fine- sreeggranddd mixture-of-exppee is also advocated by Dai et al. [30]. We analyze the effect of despen in Fig. 13 in the appendix.
[29] and
2.2 MoE Self-Attention Layers
To introduce MoE to the self-attention layers, we apply SwitchHead [31], which is an MoE method extending o-Moß 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 and NA value and output projections, which are chosen dynamically for each input. Keys and queries are computed "as usual": given an input at position f. 2, E Rodeo Ky x, W and q. = x; W where W and W E Rd x xcheet, where fl..... . H} is the head index. The expert selection for the values computed as follows:
st: = ox W sv) R
(5)
8 = arg (sp,(st KA) S (1,...,
(6)
where W E AA is the selection weight for the value and KA is the number of simultancously active experts per head, set to KΛ = in all of our experiments. The selection for the values and outputs are independent. The selection of the output is computed analo- gously using a different weight matrix W E dh NN. SO. aa, oa, (o) E N and Eg (xx) = arg topk( KA)= C {l,.... NA). Then the output y E R is calculated as follows:
h Σ
&alela, Who E Rdwa
EEE) (m))
a Attention (g.
K.. ViA e ER
(8)
yt Σ eeE
V.
where whe E dhaa and O E &Rdhaal dawca are head h, expert weight matrices for value and output respectively; st [t] sh [c] E IR are scores of expert for head h at position for value and output MoE respectively; and Attention denotes the standard softmax scaled dot attention [1] with K 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 SwitchHead, which uses no regularization, we apply the same entropy regularization we use in the feedforward layer (Eq. 4) with regularization coefficient d. (The same value is used for both value and output.)
kf)), V.a = (of...., v.),,," T'xdhar e.g., for the auto-regressive setting.
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
3

---

Layer B
Layer A
Layer
Layer B
Layomomm
Layer A
Layer B
Layer
G-2
A
Layemorm
Lspersers
Layer grouping: 8 layers with group size of 2.
Figure I:
Figure 2: The residual grows in pre-layernorm transformers.
Figure 3: MoEUT block with no layernorms the residual.
detrimental. Second, the total number of attention heads is kept relatively low, which might not be sufficient for large model. Increasing their number similarly prohibitively expensive.
Our solution to these problems is to stack multiple layers with non-shared weights to form what we call group of layers, reducing the number of experts in each o-MoE while increasing the total number of attention heads. The final network is obtained by recurrently stacking such groups that share the shared "layer" in the UT). Fig. I 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 or 9.
same parameters (in sense, redefining the group
As further observations in favor of the potential inductive bias introduced by such grouping, note that 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 complex synthetic tasks, such as ListOps. Both these observations indicate that the adjacent layers in Transformers often perform different sub- oobbooprations for single high-level step of computation that spans multiple layers. This is well aligned with our proposed grouping.
two consecutive layers to perform single operation for relatively
2.4 Novel LayerNorm Scheme for Improved Signal Propagation in Universal Transformers
Virtually all modem 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 erucial for training deep models. This indicates that the norm of the residual vector should grow as we go deeper in the network (see Fig. 2 for an illustration). However, 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, they have separate parameters which can have different scalings in different layers, and this can be observed empirically (for more details, Appendix A.2). This is not the case for UTs they have single, shared layer (or in our case multiple, repeated layers; see Sec. 2.3). 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 lay-layem, and with the fact that the original UT [10] is trained with post-layernorm. That said, mentioned above, post-layernorm also has its own limitation in terms of gradient flow [34].

---

18
Baseline MoEUT
18
Baseline
Perplexlly
16
Perplexitty
16
McEUT
e MMM
o MoE
14
14
12-
12
200
400
600
800
1000
2
4
Number of model parameters (M)
Total training MACs
x101
(a) Scaling in the number of parameters.
(b) Scaling in the number of MACs used for training.
Figure 4: Scaling of different models on C4 (with perplexity measured on a held-out subset of C4). (a) MoEUT slightly outperforms parameter-mateee matched models with no layer sharing. The gap grows with scale. (b) Given equal amounts of compute, MoEUT outperforms other models by large margin.
Here we propose an alternative method to avoid the aforementioned problems: we do not use layernomss 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 o-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 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 laye-laa,,,,,m, which positions layernorm "around" (but not on) the residual connections.
Main Experimental Results
We present main experimental results the performance and efficiency of MoEUT language modeling using the popular C4 dataset [39]. To demonstrate the versatility of our model, we also show main results 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, Following prior work [27, 31J, we measure the compute requirements in terms of the number of multiply-accumulate (MAC) operations needed in the forward pass.
refer much previous work (e.g.. [10, 15, 14, 16, 38]).
Because our models are fully MoE, they decouple the number of parameters, compute and memory requirements, and different model dimensions such dmodel and dr, number of layers. Thus, they provide greater ficxibility for model designers. We follow a simple procedure for setting the model's hyperparameters, described below. All our models use RoPE positional encodings [43] with PyTorch's 's fast attention implementation. The baseline models are pre-layernorm Transformers. For each baseline, construct parameter-matcher MoEUT model. We set dmuddel and the number of layers Vlayeers 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 1 H of the corresponding dense model, dread is set to 20head of the corresponding dense model, and we set K^ 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 o-MoE layers, we set the expert size dexpent = 128, and K = 2dmodet dexpent This halves the MAC requirements compared the dense counterpart. We the number of experts in the feedforward block, Ng. and the number of attention experts, NA 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, = 3 for our 319M parameter model, and = 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 and context length 1024, for 10a steps. This protocol allows to perform fair comparisons between different models within our
5

---

2.80
Beseline MOEUT
14.2
2.75
Perplexity
14.0
2.70
13.8
2.65
2.55
13.4
2.50
13.2
244M
728M
AABB G-1
G-2
G-3
G-6
G-9
G-18
Number parameters
Figure 5: Performance of MoEUT compared to standard Transformer on The Stack. MoEUT outperforms standard Transformers. The gap grows with scale.
Figure 6: Perplexity of 244M MoEUT models with different layer grouping options. A small group size of G = 2 works the best, showing the advantage of layer sharing.
computational budget, and it leads to high quality models, as measured by our benchmarks. For more details, see Appendix A.4.
Scaling compared to standard Transformers. Our main scaling results are shown in Fig. 4. The y-axis shows the perplexity slightly outperforms dense models with the same number of parameters (Fig. 4a), and the gap tends to grow with scale. Additionally, baseline has the same shape of feedforward layers (dmodel. K, dexpent) to our layer-shared MoEUT, but uses attention experts to keep the proportion of the attention weights close MoEUT as possible, and it also uses our peri-layernorm scheme (Sec. 2.4). We add this baseline as the model that is close to shared-layer model possible. This model performs significantly worse than MoEUT, demonstrating the clear advantage of the shared layers. Additionally, Fig. 4b 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.
held-out subset of C4. The plot shows that our MoEUT model
compare to the non-shared oMMoE model [28]. This o-MoE
Performance on code generation. To confirm the effectiveness of our model on a different task domain, here we train it subset of the "The Stack" dataset [42] which is 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 Fig. 5, and they in line with our findings the natural language domain: MoEUT outperforms the baseline.
Zero-shot performance downstream tasks. Here we evaluate the zero-shot performance of our models on six different downstream tasks: LAMBADA [44], BLiMP [45], Children's is Book Test (CBT) [46], HellaSwag [47], PIQA [48], and ARC-E [49]. For LAMBADA, we use the detokenized version from OpenAl, and we evaluate the top-1 accuracy of the last word (it can span multiple tokens; here
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 Tab. 1. We observe that our models and the baselines typically perform very similarly. MoEUT often outperforms the baseline, but the differences are marginal all cases. This confirms that our models are indeed very capable compared to standard language models. We confirm this peS2o and SlimPajama well.
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 standard 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-layernom and does not use layer grouping. Unlike ours, Adaptive Computation Time (ACT) is used the layer dimension.
We took the original code released by Tan et al. [38] and ported it to our training pipeline for fair comparison. As for MoEUT, we roughly match the model's dimensionalities and number of active channels to our dense basclines. We hyperparameter optimization for the regularization losses,

---

Table 1: Zero-shot downstream performance and perplexity on various language modeling datasets. MoEUT marginally outperforms standard Transformers in most tasks, confirming that MoEUT is indeed capable language model.
Datasct
#params
Model
PPL 4 LAMBADA f
BLIMP 1
CBT t
HellaSwag
PIQA +
ARC-E 1
Average
Baseline
18.97
21.9%
73.5%
81.3%
28.3%
59.9%
31.7%
49.4%
44M
MoEUT
18.30
23.2%
78.2%
81.1%
29.2%
61.3%
33.5%
51.1%
Raseline
14.97
28.5%
77.0%
84.4%
31.7%
62.7%
35.2%
53.2%
126M
MoEUT
14.76
27.2%
32.3%
35.3%
53.8%
C4
244M
MoEUT
Baseline
13.40
13.24
33.1%
30.6%
79.7%
78.5%
85.3%
86.0%
34.5%
35.7%
65.2%
64.9%
36,9%
36.4%
55.5%
55.6%
Baseline
12.81
78.5%
87.2%
36.1%%
37.2%
56.6%
319M
MoEUT
12.65
30.8%
80.2%
86.9%
37.3%
67.0%
37.3%
56.6%
Baseline MoEUT
11.59
37.8%
80.7%
88 2%
40.5%
67.7%
39.3%
59.0%
11.34
36.0%
80.8%
58.4%
41.8%
69.2%
39.6%
59.3%
Baseline MoEUT
11.15
38.4%
81.2%
89.0%
42.0%
68.6%
39.7%
39.8%
10HM
10.90
60.7%
44M
Baseline MoEUT
11.46
13.2%
66.5%:
68.6%
28.5%
56.3%
32.0%
44.2%
11.09
13.1%
68.7%
69.6%
28.3%
55.1%
31.4%
44.4%
PeS2o
Baseline MoEUT
18.7%
72.8%
78.0%
30.4%
35.0%%
48.5%
244M
8.52
19.4%
73.5% %
77.4%
30. 1%6
56.3%
38.6%
48.7%
Baseline
16.42
20.0%
72.8%
80.7%
27.5%
$7.0%
31.6%
48.3%
44M
MoEUT
15.77
19.8%
75.9%
82.1 %
28.00%
57.5%
32.1%
49.2%
SlimPajama
Baseline MoEUT
11.51
31.9%
78.6%
87.3%
31.7%
60.9%
36.6%
54.5%
11.47
30.7%
80.2%
86.88%
32.0%
61.7%
35.8%
54.5%
Baseline 9.56
38.8%
80.5%
39.9%
37.6%
64.5%
38.7%
38.3%
1040M
MoEUT
82.5%
90.2%
58.7%
and we found that a minimal regularization is necessary for stabilizing the training. However, larger regularization tends hurt performance significantly. All other hyperparameters set based Tan et al. 's biggest translation experiments. The results are shown in Fig. 7. Effectively, SUTs, which lack our specific methods, have a significant performance disadvantage compared MoEUT and the sarameter-matched dense baseline. Upon careful investigation, we found that most of this poor performance comes from the ACT mechanism that the authors advertise as one of the main components of their model. After removing the ACT, the performance improves dramatically. However, even with this setup, it underperforms both MoEUT and the standard Transformer baseline. This is also confirmed on downstream tasks in Tab. 2 in the appendix. Moreover, as we show in Appendix A.7, our model runs much faster and uses only a fraction of the memory required for the SUT. To the best of our knowledge, we are not aware of any prior UT architectures that are both competitive and efficient in language modeling.
our
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 non-shared layers within the group. G = corresponds to the model used all other analyses. G 1 fully shared-layer model, without any grouping, and G = 18 corresponds to the baseline fully non-shared a-MoF model [28]. All hyperparameters are identical among all models, except for the number of MLP experts (Ng) and attention experts (Na). which are adjusted to match the parameter count of the dense baseline. In Fig. 6, we observe that = 2 optimal, and the recurrence the layer dimension is indeed beneficial. Another interesting question is whether the grouping described in Sec. 2.3 and Fig. is the right way to stack layers. Let us call the two layers in the group and B. The grouping we discussed so far stacks layers in the form of "ABABAB", e.g., for 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. can be seen that the grouping proposed in Sec. 2.3 indeed works significantly better. In fact, the AABB-style stacking is almost bad 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 7 719M- -arammetet
7

---

40-
Baseine MaFUT
SUT w.o. ACT
19-
PostLN PreLN
SUT
Perplexiy
Replexx
PeriLN
20-
10
13
C4. 44M
C4, 244M
peSoo, 44M
peS2o, 244M
44M
244M
Figure 7: Comparing MoEUT and SUT [38]. MoEUT outperforms both the original SUT and our improved version by large margin.
Figure 8: Comparing layernorm variants. Peri-layernorm outperforms both pre- and post-layemnomm
model, the gap between peri-norm and post-norm is marginal: 11.29 perplexity points compared 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.
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. 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 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 lay- ers. Note that MME-bassed shared-layer models could, in theory. assign different experts to different lay- ers to "emulate" regular non-shared models. If this were the case, the model would resemble a regular Transformer instead of Universal one. To confirm that our model is more versatile than that, we ana- lyze whether certain experts in the MLP layers are activated only in specific layers. We measure how many times each expert is activated in each layer. This allows us 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 expert activations in that layer. Fig. 9 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 a wide range of layers where the expert is activated. Experts seem to be active in continuous sequence of lay- ers. This can be seen in the wide, vertically lined structure. We can conclude that MoEUT is capable of specializing in specific layer if necessary and sharing weights between them when advantageous.
Per-token expert selection diversity. Here we analyze the diversity of expert selection in the MLP layers for 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. 10. On the x-axis, the tokens ordered differently for each layer based on the number of experts used 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 fully "specialized"" to consistently we look at the first layer (blue curve): the number of unique experts used is low, i.c., the selection mainly depends on the token's '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 tokens/layers, refer to Appendix A.6.
single set of K experts. This is almost the case for many tokens
Expert selection dynamics in individual columns/positions. have been cumulative statistics in different input sequences and positions. We might
So far, the results

---

5
0.75
#Experts Used
60
Layer
Layer
0.50
Layer
11
Layer
13
0.25
Layer 13
17
0.00
20
Layer
100
200
300
Expert ID
0
200
400
600
800
1000
Token
Figure 9: Layer preference of different experts. Most experts are used in multiple layers, while some them (see, bottom right) specialize to certain layers, showing the flexibility of model.
Figure 10: No. of unique experts used in differ- ent layers. Tokens are routed to many different experts (depending on the context), especially in the middle layers.
wonder about the selection behavior for the individual Transformer columns." 3 Is ex- pert 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. 11. There is non-negligible overlap be- tween the selected experts in subsequent layers; however, is far from complete overlap. This indicates that experts usually change dynamically single column, performing different functionality in different layers.
1.00
5-
0.75
0.50
0.25
17
13
17
Layer
Overall, our analysis suggests that MoEUT capable of dynam- ically adapting its expert selection mechanism to 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 11: Instance-level aver- age expert selection similarity between layers. Individual to- kens routed to diverse of experts across the layers.
Discussion and Limitations
More background on UT. We emphasize that our focus is on developing scalable and performant Universal Transformers for parameter-dominatiee tasks such language modeling. This has been long-standing limitation 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 achieve compute--fficiency-l contrast, our MoE-based approach is more compute-efficien 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 them have only focused on applying MoE the feedforward layers. Notable exceptions are mixture-of-atteniin [27] and SwitchHead [31] (Sec. 2.2), which focus on MoE self-attention layers. In addition, until recently, it has been considered challenging make MoE-based language models competitive against the dense baseline in the
matched setting settings). MoEUT, we use o-MoE [28], an MoE design that has been shown to be competitive even in such a setting.
(unlike FLOPs/MAC-matee
By representing the Transformer's activations in 2D grid, with token positions on the x-axis and depth on the axiis, "columns" correspond to all hidden activations across depth given token position.

---

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 with more steps perform better. Sometimes, layer grouping is also used to up-scaling pretrained models [55].
Our current implementation of the MoE layers uses the Triton kernel released with o-MoE [31] for both the attention and the MLP parts of the model. This implementation known to be suboptimal [31]. Compared to the standard Transformer with FlashAttention [56], MoEUT model trains 1.5-2x slower. We estimate that with more optimal implementation, the training speed should be close to the dense model, while inference should run faster.
simitation/mmplementatoonn
Massive scaling. Our experiments used 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 very large compute cluster, but the costs could also be mitigated somewhat by work optimizing our CUDA kernel.
Conclusion
We present MoEUT, novel coxture-of-Expertbbassee 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- models. Our MoEUT allows for training competitive UTs on parameter-done -ommnnttdd tasks such as language modeling, while being significantly less compute intensive than the baselines without layer sharing. We break this long standing limitation of UTss 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 helps revive research interest Universal Transformers at larger scales, and serves as 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 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 IDSIA for providing part of the compute used for this project even after the authors left the lab.
References
[1] 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.
[2] Jürgen Schmidhuber. Learning to control fast-weight memories: An alternative to recurrent nets. Neural Computation, 4(1):131-13, 1992.
[3] Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.
[4] Tom Brown al. Language models are few-shot learners. In Proc, Advances in Neural Information Processing Systems (NeurIPS), Virtual only, December 2020.
10

---

[5] OpenAl. ChatGPT. https://openai: http:///i...ooomoooo//tttggae/e/ 2022.
[6] OpenAl. GPT-4 technical report. Preprint arXivv23033.08 08774, 2023.
[7] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timo- thé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.
[8] 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.
[9] 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.
[10] Mostafa Dehghani, Stephan Gouws, Oriol Vinyals. Jakob Uszkoreit, and Lukasz Kaiser. Univer- sal Transformers. In Int. Conf. on Learning Representations (ICLR), New Orleans, LA, USA, May 2019.
[11] Jeffrey L. Elman. Finding structure in time. Cognitive Science, (4(2):19--211 1990.
[12] Michael 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.
[13] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, pages 1735-1780, 1997.
[14] 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.
[15] Róbert Csordás, Kazuki Irie, and Jürgen Schmidhuber. The devil in the detail: Simple tricks improve systematic generalization of Transformers. In Proc. Conf. on Empirical Methods in Natural Language Prucessing (EMNLP), Punta Cana, Dominican Republic, November 2021.
[16] 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.
[17] Jürgen Schmidhuber. Self-delimiting neural networks. Preprint arXiv: perriiv:220...0 1210.110100011 18, 2012.
[18] Alex Graves. Adaptive computation time for recurrent neural networks. In Int. Conf. on Learning Representations (ICLR) Workshop Track, Vancouver, Canada, April 2016.
[19] 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.00 003600,0...11..... 2020.
[20] 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.
[21] Robert A. Jacobs, Michael I. Jordan, Steven J. Nowlan, and Geoffrey E. Hinton. Adaptive mixtures of local experts. Neural Compututaioon 3(1):79-88 87, 1991.
[22] John B. Hampshire II and Alexander H. Waibel. The meta-pi network: connectionist rapid adaptation for nhigh-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.
[23] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe 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.
11

---

[24]
Dmitry Lepikhin, HyoukJoong Lee. Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. GShard: Scaling giant models with condi- tional computation and automatic sharding. In Int. Conf. on Learning Representations (ICIR), Virtual only, May 2021.
[25]
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Preprint arXiv:2101.0396) 2021.
[26]
Aidan Clark, Diego de Las Casas, Aurelia Guy, Arthur Mensch, Michela Pagannni, 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.n 01 169, 2022.
[27] 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 4160-4162, Abu Dhabi, United Arab Emirates, December 2022.
[28] 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 2023.
[29] Jakub Krajewski, Jan Ludziejewski, Kamil Adamczewski, Maciej Pióro, Michal Krutul, Szy- mon Antoniak, Kamil Ciebiera, Krystian Król, Tomasz Odrzygózdz, Piotr Sankowski, Marek Cygan, and Sebastian Jaszezur. Scaling laws for fine- fine-grained mixture of experts. Preprint D.Xiv:2402.07871 2024.
[30] 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
0tt::2200.000666 2024.
[31] Róbert Csordás, Piotr Piękos, Kazuki Irie, and Jürgen Schmidhuber. SwitchHead: Accelerating transformers with mixture-ofeexperts attention. In Preprint arXiv:2312.07798 December 2023,
[32] 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, 7ac 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 waniiinngaand-duuutiiinn Circuits Thread, 2022. https:/transsforme-irrcut httpa//mmmmmiiiiuuuuuooooo/ooooo heads/index.html.
[33] 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.
[34] 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.
[35] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. Preprint arXiv: 1607.06450, 2016.
[36] Simon J. Thorpe. Local vs. distributed coding, Intellectica, 8:3-40, 1989.
[37] 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.
[38] Shawn Tan, Yikang Shen, Zhenfang Chen, Aaron C. Courville, and Chuang Gan. Sparse universal transformer. In Proc. Conf. on Empirical Methods Natural Language Processing (EMNLP), pages 169-179, Singapore, December 2023.
12

---

[39] 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 (JMLRR, 11:1140111140:: 2020.
[40]
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. ct/
URL hal.:tttps:/huugginggaae
ama-627B.
[41]
Luca Soldaini and Kyle Lo. peS2o (Pretraining Efficiently on S2ORC) Dataset. Technical report, Allen Institute for AI, 2023. attpss///ittan
[42]
Denis Kocetkov, Raymond Li, Loubna Ben Allal, Jia Li, Chenghao Mou, Carlos Muñoz Fer- randis, 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.
2022.
[43]
Jianlin Su, Yu Lu, Shengfeng Pan, Bo Wen, and Yunfeng Liu. RoFormer: Enhanced transformer with rotary position embedding. Preprint arXiv:2104.099
2021.
[44] 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 Compu- tational Linguisties (ACL), Berlin, Germany, August 2016.
[45] Alex Warstadt, Alicia Parrish, Haokun Liu, Anhad Mohananey, Wei Peng, Sheng-Fu 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.
[46] 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.
[47] 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 18000 Florence, Italy, August 2019.
[48] 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.
[49] Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try ARC, the A12 reasoning challenge. Preprint 1rXivv1803.000
2018.
[50] 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.
[51] Jacob Devlin, Ming- Bing-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 1-166,,,,,,,1 Minneapolis, MN, USA, June 2019.
[52] 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.
[53] 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. I 44802 2023.
[54] 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.
[55] Dahyun Kim, Chanjun Park, Sanghoon Kim, Wonsung Lee, Wonho Song, Yunsu Kim, Hyeon- woo Kim, Yungi Kim, Hyeonju Lee, Jihoo Kim, Changbae Ahn, Seonghoon Yang. Sukyung Lee, Hyunbyung Park, Gyoungjin Gim, Mikyoung Cha, Hwalsuk Lec, and Sunghun Kim. SOLAR
13

---

10.7B: Scaling large language models with simple yet effective depth up-scaling, Preprint arXiv:3312.1516,, 2023.
[56] Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Attentot:: Fast and memory-efficient exact attention with IO-awareness. In Proc. Advances in Neural Information Processing Systems (NeurIPS), New Orleans, Louisiana, USA, December 2022.
[57] 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 Vito, 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 (NeurlPS), pages 8024-8035, Vancouver, Canada, December 2019.
[58] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In Int. Conf. on Learning Representations (ICLR), New Orleans, LA, USA, May 2019.
[59] 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
14

---

A Appendix
A.1 Broader Impact Statement
We consider this work to be 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 potentially be a jump towards Artificial General Intelli- gence, which might have unforeseeable consequences, both positive and negative. Additionally, 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 beneficial because of the reduction in energy usage, but it might also enable casier generation of harmful content like fake news or deepfakes.
A.2 Growing Residual Norm Standard Transformers
In Sec. 2.4, we discussed the issue of the growing residual norm in the standard Transformers. Here, we measure the L norm the difference the residual before and after applying a standard Transformer layer (both the attention and the MLP block) in different layers of 44M parameter Transformer trained on C4. The results are visualized in Fig. 12. can be seen that the norm of the updates indeed grow in later layers.
Update magnitude
100
10
15
Layer
The update magnitude of different layers in 44M parameter Transformer on C4. The norm of the updates grows throughout the layers compensate for the residual growth (see Sec. 2.3 for more details).
Figure 12:
A.3 Zero-Shot Downstream Performance of SUT-variants
In addition evaluating the perplexity of different SUT variants in Fig. 7, we also show their zero-shot downstream performance in Tab. 2. can be seen that MoEUT consistently outperforms both SUT and SUT without ACT.
Table 2: Zero-shot downstream performance and perplexity of different SUT variants compared to MoEUT and the unshared baseline. MoEUT outperforms both SUT variants.
Dataset
fparams
Model
PPL.
LAMBADA
BLAMP 1 CBT 1
HellaSwag
PIQA 1 ARC-E.
Average
Baseline
18.97
21.9%
73.5%
81.3%
28.3%
59.9%
31.7%
49.4%
81.1%
44M
SUT
40.50
1.2%
65.3%
51.1%
26.4%
57.8%
31.9%
39.0%
C4
SUT w.s. ACT
21.51
18.1%
72.8%
66.3%
27.5%
59.1%
32.5%
46.0%
Baseline MoEUT
13.40
33.1%
78.5%
86.0%
34.5%
64.9%
36.9%
$5.6 %
244M
SUT
20.05
20.5%
71.0%
68.5%
28.29
60. 60.1%%
32.7%
46.8%
SUT w.o. ACT
14.58
27.8%
77.0%
75.9%
32.7%
2%
52.0%
Baseline
11.46
13.2%
66.5%
68.6%
28.5%
56,
32.0%
44.2%
44M
SUT
25.04
0.5%
59.25
38.1%
26.25
55 .0%
31.1%
35.00%
w.o. ACT
12.68
11.7%
66.5%
53.9%
28.0%
56. 1%
31.5%
41.3%
PoS2o
Buseline
8.55
18.7%
72.8%
78.0%
30.4%
56.3%
35.0%
48.5 59.5%
MoEUT
8.52
19.4%
73.5%
77. 47.4%%
30.1%
56.3%
35.6%
48.7%
244M
SUT
20.44
0.5%%
60.9%
42,8%
26.7%
55.3%
33.0%%
36.5%
w.o. ACT
16.8%
64.8%
15

---

A.4 Hyperparameters
All our models are trained in PyTorch [57] with a batch size of 64, context length of 1024, for 100k iterations, learning rate of 0.00025, AdamW optimizer [58] with default hyperparameters. weight decay of 0.01. They are trained on single node in data-parallel manner. The learning rate decayed to 10% of its initial value using cosine decay. We use gradient clipping of r. and warmupp linear learning rate warmup steps (see Tab. 3). None of our models uses dropout. For the entropy regularization of the MLP expex selection, we use ? 0.01 and for SwitchHead attention 5 = 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 the given dataset. All models are trained with mixed precision. The hyperparameters of the SUT models can be found in Tab. 4. Note that the meanings of the parameters are not directly analogous to ours. Please refer to Tan al. [38] for more details.
Table 3: Hyperparameters of different models used in our main experiments.
Model
#params
Taayers
G
daxide
dr
H
NA
dhed
N
K
Nu
Baseline
45M
16
412
2053
10
41
0.1
MoEUT
44M
16
82
155
o-MoE
44M
16
16
412
82
17
12
0.1
Baseline
126M
16
3072
16
48
4000
0.25
MoEUT
126M
18
768
10
96
254
12
4000
0.25
o-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
1024
4
10
128
387
16
4000
0.25
o-MoE
244M
18
18
1024
128
4000
0.25
Baseline
319M
24
1024
4110
16
64
4000
0.25
MoEUT
318M
24
1024
10
338
4000
0.25
o-MoE
320M
24
24
1024
128
16
4000
0.25
Baseline
729M
36
1280
5120
20
64
0.25
MoEUT
727M
36
1280
13
128
467
20
4000
0.25
o-MoE
731M
36
36
1280
128
50
20
4000
0.25
Baseline
1044M
36
1536
6144
24
64
4000
0.25
MoEUT
1040M
36
1536
12
128
565
24
4000
0.25
Table 4: Hyperparameters of SUT models used in our experiments.
#params
Maayers
dmodel
dexpest
H
NA
danee
dhead
Ng
K
Csum
Lacr
Nwarmap
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.5
The Effects of dexpert and K
Here, we analyze the effect of the expert size given fixed amount of compute (dexpen K being kept constant). The results are shown in Fig. 13. It can seen that using fine-grained mixture of experts (small experts) is indeed critical for good performance. In models, smaller experts significantly decreases the compute efficiency 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.
dexpert 128. Using
We also analyzed the performance of our MoEUT model with different numbers of active experts the MLP layer. This varies the amount of compute spent the layer, and the number active channels. We show the results in Fig. 14. Increasing the number of active experts always improves performance, the returns diminish. We chose 16 for our experiments because of efficiency reasons.
16

---

13.7
Perplxiit
128
258
512
1024
douper
Figure 13: Performance of our 244M MoEUT on held-out subset of C4 with expert sizes (despen) in the MLP layer. The smallest expert size performs the best.
13.5
Perpleaily
13.4
13.2
13.1
16
Figure 14: Performance of our 244M MoEUT on a held-out subset of C4 with different number of active experts for the o-MoE layer. Increasing the number of experts always helps, but the returns diminish.
A.6 Additional Analysis
Expert specialization to tokens/layers. Now conversely to the "per-token expert selection diversity" analysis presented the main text, we analyze whether the experts activated by a token are layer- specific for that specific token. For this, count the number of unique experts used by each token and compute the corresponding proportion for 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 observe that for the most frequent tokens (toward the left part of the plot), high scores near 1.0 are obtained in multiple layers for a given token; this means that (almost) all experts used by that specific token are used in multiple layers. In contrast, for the less popular tokens (toward the right part of the plot). In addition, the set of experts selected in the early layers is typically less diverse than for the rest of the layers: only small fraction of the used experts are present there. This consistent with the findings for Layer in Fig. 10.
observe that experts tend to be more layer-specific
1.00
0.75
Layer
0.50
13
0.25
1000
2000
3000
4000
5000
Token
Figure 15: Proportion of experts used in a specific layer out of all unique experts used by given token. On the x-axis, the tokens are ordered by decreasing frequency of occurrence. The less frequent a tokens is, the more layer-specific are the experts used by that token.
A.7
Wall Clock Time and Memory Comparison
To show the real-world resource usage our models directly, we run a controlled experiment on identical hardware with our 244M parameter model and the corresponding baselines.
measured the training iteration time and memory usage on V100 32GB GPUs. Here one "iteration" corresponds to effective batch size of 64x 1024 tokens for all models. The training
17

---

iteration time was measured by using batch size for each model that fits GPUs; models require either 1 or 2 gradient accumulation steps to achieve the effective batch size depending on their memory requirement. We measured the training time right after initialization and a warmup period. The memory usage is measured using grad accumulation steps for all models for fair comparison. Note that around 3GB of memory is used by the model parameters and optimizer state on each GPU. We show the results in Tab.
Even though our MoEUT with the current kernel implementation is slower than the corresponding dense non-shared layer transformer, is significantly faster and uses much less memory compared to alternative UT variants.
Table Wall-Clock Time for the waraadd pass and the total memory usage our training loop with different 244M parameter models 8 V100 GPUs, which batch size of 64x 1024 tokens. MoEUT 1.7x slower with our suboptimal MoE kernel implementation than the standard transformer, but it is much faster than the other UT variants. It also uses much less memory, allowing training larger scales.
Model
ms/batch
Memory usage
Non-shared Transformer
443
9.2 G
Naive UT
3559
25.9
MoEUT
772
9.0
SUT
1344
23.4
A.8 Compute Requirements
We report the hardware used and the required wall clock time for our main experiments in the paper in Tab. All experiments are performed private clusters. The duration reported "hh:mm" format. We report the number of GPUs (Naru) used for that specific experiment (and nor the total number of GPUs in the system). For the number of CPUs (Ncpu) 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 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 corresponding 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.
18

---

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

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
We summarized the motivation, method, and main findings in these sections.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We discuss the speed issues of our current implementation and that our experiments do not consider algorithmic tasks and generalization.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
20

---

Justification:
Our paper is empirical.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: show all the hyperparameter configurations Appendix A.4, and we provide the code for our experiments.
Guidelines:
The answer NA means that the paper does include experiments. If the paper includes experiments, No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
21

---

Answer: [Yes]
Justification:
We provide the code for our experiments. auto-downloads all the data that
needs.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run reproduce the results. See the NeurlPS code and data submission guidelines (https: /Ct///ooodeessseeeeeeeeee ssionnolicyyyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: The method of choosing the hyperparameters is described in Sec. 3, and the full table of hyperparameters is presented in Appendix A.4.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: Our experiments involve large models that are very expensive to train, and we do not have sufficient compute resources run multiple seeds them.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by emror bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
22

---

The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We report the type of hardware used for our main experiments in Appendix A.8. We also give rough estimate of the total amount of resources used.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https:/ http:::////uiipp// cctpublic//tttoecoeeeceeeeeeeett ipppbbic//ttcccciiiddiieees
Answer: [Yes]
Justification: We read the Ethics guidelines, and to the best of our knowledge, we are complying with it.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there is a special consid- eration due to laws or regulations their jurisdiction).
Broader Impacts
Question: Does the paper discuss both potential positive socictal impacts and negative socictal impacts of the work performed?
Answer: [Yes]
Justification: We consider our paper be a foundational research paper without direct consequences. Despite this we also discuss the potential consequences in Appendix A.1.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
23

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied partieular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (c.g., gated release models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of MI).
11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification: The models in this paper are small by modern standards and we do not release pre-trained weights.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (c.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: Our code is under MIT license and the paper is CC-BY 4.0. To the best of our knowledge, we always credit the reused code if we reuse any.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: We open source our code and provide instructions on how to run it. Upon acceptance we plan also provide an easy to use, well documented single-file version of our model.
Guidelines:
The answer NA means that the paper does release assets.
Researchers should communicate the details the dataset/code/mooell as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: do not work with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
We do not work with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
25

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
26


Reviews:
[{"review_id": "GteDwo0nlM", "reviewer": "Reviewer_83mA", "summary": "The paper suggests to use Sparse MoEs together with Universal Transformers to overcome the parameter-count limitation that the latter have when parameters are shared over consecutive layers. In particular, the work suggests to use $\\sigma$-MoEs (that use sigmoid activation function in the router, rather than the more popular softmax activation), and use fine-grained experts (selecting many small experts over a large pool, rather than one or two large ones over a smaller pool). The goal is to make Universal Transformers competitive on language modeling tasks, where they haven’t excelled in the past (presumably due to the smaller parameter count, compared to dense models of the same time complexity).\n\nTwo important novelties that are introduced are Layer Grouping in UTs (this allows to chose the degree of parameter sharing across layers in the UT), and relocating the LayerNorm operations (called PeriLayerNorm, as opposed to pre and post LayerNorm).\n\nThe paper includes experiments on several pretraining tasks (C4, peS2o, SlimPajama) and zero-shot evaluation on several downstream tasks. In both cases, the results achieved by MoEUT are competitive with a standard dense transformer with the same number of parameters, and often slightly better.", "questions": "- In figure 10 and 11, do the expert indices correspond to the experts in the attention or in the feed-forward part of the block? Do the trends differ across the two?", "limitations": "No negative societal impact particular to this work, in my opinion.", "rating": 7, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The presentation of the work is excellent. From the motivation, to the explanation of the proposed method, related works, and the description of the experiments.\n- The main experiments show that the goal of making UTs competitive with standard dense transformers in language modelling tasks was achieved.\n- Many additional ablation experiments were conducted to try to explain how experts are selected across different layers, and tokens.\n- An effort was put on implementing reasonable baselines such as $\\sigma$-MoEs and Sparse UTs.\n- The paper clearly states current limitations (e.g. not the most efficient implementation, which results in experiments that are 2x slower than the dense counterpart).", "weaknesses": "- The proposed method matches dense transformers on language modeling, but it’s barely better. This begs the question: why using MoEUTs then, rather than the simpler baseline? The answer could be better MAC efficiency at the same memory cost, but this is not clearly represented in Table 1 (and as the authors point out, it’s actually slower due to implementation limitations).\n- The perplexity reported for Sparse UTs is alarmingly high. The paper mentions that proper hyperparameter tuning was done, but I’m wondering if the authors reached out to the SUT authors to make sure that everything was implemented correctly.\n- A similar thing happens with $\\sigma$-MoEs, which performs worse than the dense baseline when matching the parameter count, and most troubling barely better when matching MACs.\n\n\n-------\nEdit after rebuttal: The authors addressed most of my concerns in the rebuttal. Given their response to my and to other reviewers' comments, I'm increasing my score to \"Accept\".", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their insightful review and for the positive comments on the clarity and methodology of the paper. Please find our responses as follows:\n\n> The proposed method matches dense transformers on language modeling, but it’s barely better. This begs the question: why using MoEUTs then, rather than the simpler baseline? \n\nOur goal is to pave the way towards a general model that can be used for LM but is also good at systematic generalization, eventually arriving at foundation models that are more trustworthy and data efficient. In this paper, we address a single step toward that goal: we present a way to overcome the long-standing efficiency limitations of UTs (which are known to generalize better than standard transformers), which have prevented their scaling to sizes required for LLMs. \n\n> it’s actually slower due to implementation limitations\n\nIt is true that our current proof-of-concept kernel is slower than a standard transformer despite requiring fewer flops, but we believe this can be mostly mitigated by better kernels. Even with the current kernels, our method is significantly faster than a naive UT or the next best UT baseline, SUT, and for most of our experiments naive UTs are not viable to run because of their memory usage and speed (see time/memory usage measurements below). Rather than replacing Transformers, we are advocating for more research on Universal Transformers with our method for larger scales.\n\nNote additionally that we investigated these models in a parameter-matched setting because we were interested in comparing their expressiveness. The number of parameters in our MoEUT can be increased very cheaply (as with other MoEs), and it is possible to achieve significantly better perplexity for marginally higher compute, potentially justifying the wall-clock time slowdown even with the current kernels. (We use a parameter-matched setup instead of this because this would give an unfair advantage to our MoEUTs compared to the baselines).\n\nSpeed/memory usage of different models, showing that MoEUT is much faster compared to naive UT and SUT and uses less memory:\n| Model | Time / training step | Memory usage/GPU |\n| ---- | ---- | ---- |\n| Dense |  443 ms/iter | 9.2 Gb |\n| Naive UT |  3559 ms/iter | 25.9 Gb |\n| MoEUT | 772 ms/iter | 9.0 Gb |\n| SUT |  1344 ms/iter   | 23.4 Gb |\n\nWe measured the training iteration time and memory usage on 8 V100 32Gb GPUs. Here one “iteration” corresponds to an effective batch size of 64x1024 tokens for all models. The training iteration time was measured by using a batch size for each model that fits GPUs; models require either 1 or 2 gradient accumulation steps to achieve the effective batch size depending on their memory requirement. We measured the training time right after initialization and a warmup period. The memory usage is measured using 2 grad accumulation steps for all models for a fair comparison. Note that around ~3Gb of memory is used by the model parameters and optimizer state on each GPU. \n\n> The perplexity reported for Sparse UTs is alarmingly high\n\nWe confirmed via personal communication through a colleague that SUT does not work well for LM. We used the official code of the authors, with minimal modification required to adapt it to our codebase. Additionally, we conducted some ablations on the SUT models, and we found that the main cause of the bad performance is the ACT used by their model (which was presented as a fundamental building block of SUT). By disabling ACT, the gap between MoEUT and SUT is much smaller, while MoEUT remains consistently better:\n\n| Dataset | Model Size | Baseline | MoEUT | SUT w.o. ACT | SUT | \n| --- | --- | --- | --- | --- | --- |\n|C4| 44M | 18.9 | 18.2 | 21.5 | 40.5 |\n| | 244 M | 13.3 | 13.2 | 14.5 | 20.0 | \n| peS2o |  44M | 11.5 | 11.1 | 12.7 | 25.0 |\n| | 244M | 8.6 | 8.5 | 9.3 | 20.4|\n\n> A similar thing happens with 𝜎-MoEs, which performs worse than the dense baseline when matching the parameter count.\n\nPlease note that our baselines are much stronger than those reported in the 𝜎-MoE paper: on C4, we achieve a perplexity of 13.4 using 244M parameters vs. 17.79 reported by the 𝜎-MoE paper using 266M.\n\nThis difference comes from two modifications. First, the 𝜎-MoE paper follows the experimental protocol of Transformer XL: we used their official 𝜎-MoE codebase, but improved their baseline by using RoPE and no XL cache. Second, they use dropout in the FFN layers of their baseline, and ‘expert dropout’ in the 𝜎-MoE. Here we disabled all dropouts in all our models as we use sub-epoch training. This resulted in perplexity improvements with a higher gain for the baseline than for 𝜎-MoE. \n\n> In figure 10 and 11, do the expert indices correspond to the experts in the attention or in the feed-forward part of the block?\n\nThose expert indices correspond to the feed-forward part. We will clarify this in the final version. Thank you for pointing this out!\n\nWe believe our response above resolves all the concerns that the reviewer has raised. If the reviewer finds our response useful, please consider increasing the score. Thank you very much.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "o3YFUyX82w", "reviewer": "Reviewer_eZCx", "summary": "This paper focuses on the problem of inefficient parameter-computation ratio in Universal Transformers (UT). UT shares parameters across layers but reduces the parameter count significantly. One naive approach is to scaling up the layer size, however, cannot be easily achieved due to the prohibitive computational resource requirements. MoEUT, the proposed approach, exploits MoE architecture in both attention and feed forward layers to address this issue. To achieve the similar performance as the standard Transformers, the authors first introduce layer grouping, which allows non-shared weights among the layers within a group. Second, they remove the layernorm in the main data path in standard Transformer to avoid gradient flow and add layernormaround the residual connections. Experiments shows MoEUT outperforms standard Transformers slightly while incurring less compute and memory costs.", "questions": "It seems the number of experts is very large as a default setting in the evaluation. What if the total number of experts is just, let's say, 8 or 16 and k is 1 or 2, similar like the sparsely-activated MoE, how the performance would change?", "limitations": "Limitations are addressed in the paper and there is no negative societal impact of this work.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The idea of exploiting MoE to address the parameter-compute ration in Universal Transformers is interesting. MoE itself incorporates ideas of expert sharing among similar tokens. It is an intuitively reasonable combination.\n2. Sound explanations of the rationale of designing layer grouping and peri-layernorm are presented in the paper together with cited works to further support the claims. \n3. The evaluation and analysis are comprehensive with deep dive to the design components.", "weaknesses": "1. The claim of using significantly less compute and memory should be further supported by evaluation numbers. Though Table 4 have presented detailed training hardware information of the experiments reported in the paper, it is very hard to determine the quantitative computation and memory costs. A general number like GPU hours, GPU duty cycles, DRAM utilization or costs in terms of dollars would be a much better indicator rather than existing presentation.\n2. In Section 2.1, it is mentioned the load balancing loss is now considered for each sequence due to the loss explosion issue when applying on the entire batch. I may not fully understand the difference here and why full batch balancing loss would cause explosion while sequence-level won't. Further elaboration or evaluation could be helpful in explaining this.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for the valuable review and for many positive comments on the methodology of our paper. Please find our responses as follows:\n\n> The claim of using significantly less compute and memory should be further supported by evaluation numbers. Though Table 4 have presented detailed training hardware information of the experiments reported in the paper, it is very hard to determine the quantitative computation and memory costs.\n\nWe will update Table 4 in the final version if the paper is accepted, to include the total GPU-hours (N_GPU * duration in the current table) for better readability. Comparing memory this way is nontrivial, as it does not scale linearly with the number of GPUs, as we are using simple data-parallel training. However, here are some direct comparisons on identical hardware for the 244M parameter models that allow directly comparing them on identical hardware, showing that MoEUT is much faster and uses less memory than the alternative UT variants:\n\n| Model | Time / training step | Memory usage/GPU |\n| ---- | ---- | ---- |\n| Dense |  443 ms/iter | 9.2 Gb |\n| Naive UT |  3559 ms/iter | 25.9 Gb |\n| MoEUT | 772 ms/iter | 9.0 Gb |\n| SUT |  1344 ms/iter  | 23.4 Gb |\n\nWe measured the training iteration time and memory usage on 8 V100 32Gb GPUs. Here one “iteration” corresponds to an effective batch size of 64x1024 tokens for all models. The training iteration time was measured by using a batch size for each model that fits GPUs; models require either 1 or 2 gradient accumulation steps to achieve the effective batch size depending on their memory requirement. We measured the training time right after initialization and a warmup period. The memory usage is measured using 2 grad accumulation steps for all models for a fair comparison. Note that around ~3Gb of memory is used by the model parameters and optimizer state on each GPU.\n\n> I may not fully understand the difference here and why full batch balancing loss would cause explosion while sequence-level won't. \n\nWe cannot provide a mathematical proof here, but we do have one hypothetical but intuitive explanation. If the balancing is done at the batch level, the experts can specialize in different *sequence types*. Since the batches are sampled randomly, the “variance” of utilization rate between the active experts used in different forward passes might vary highly. Some experts may specialize in *rare* sequences, thus they will be rarely trained, resulting in a “slow expert collapse”. When they are reactivated at a much later stage of training, they may cause sudden loss explosion as their training “stage” lags behind that of other more frequently used experts. On the other hand, if the balancing is done on a sequence level, the model is encouraged to use various experts in each sequence, eliminating some of this variance. This still allows some expert specialization, because the balancing is relatively weak. \n\nUsing a significantly larger batch size would probably also stabilize the training, but it is prohibitively expensive with our resource budget. It seems like a useful next step for better-resourced teams, and so we will emphasize this in our next version.\n\n\n\n> What if the total number of experts is just, let's say, 8 or 16 and k is 1 or 2, similar like the sparsely-activated MoE, how the performance would change?\n\nThis is an excellent question. In our experiments, we found K=1 performed significantly worse than K > 1. In order to match the MACs of the dense model, the number of experts should be at least K*N_layers, otherwise the activated experts would be “wider” than the FFN of the baseline. Furthermore, to keep the number of parameters constant, d_expert should be significantly increased. \n\nWe study the effect of increasing expert size (d_expert) in Fig. 13 in the appendix, while keeping K constant. Increasing the expert size is detrimental. Decreasing K would have an additional negative effect. We analyze this independently of K in Fig 14. \n\nIf the reviewer thinks additional experiments could add value to the paper, we can run an experiment with some of our MoEUT methods in a max-d_expert configuration and K=2 to see the effects of pushing these parameters to the extreme.\n\nThese questions are extremely valuable for us and will enable us to improve our paper. Thank you very much.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "RRHgld1f6T", "reviewer": "Reviewer_ZbxN", "summary": "This paper introduces a novel application of the mixture of experts (MoE) architecture within both the MLP and attention modules of the Universal Transformer network. The integration of MoE is further complemented by an innovative sequence-level routing regularization technique, which the authors argue enhances training stability. Additionally, the paper proposes the incorporation of layer grouping and a new layer normalization scheme, aimed at boosting model performance. Compared with baselines, the proposed MoEUT can achieve better perplexity with the same number of parameters and obtain competitive performance in various downstream tasks.", "questions": "1. **Layer Grouping and Novel Layer Norm Performance Contribution:**\n   - Do the layer grouping and the novel layer normalization contribute independently to the performance improvements observed, or are these contributions specifically related to the architecture of the Mixture of Experts (MoE) network?\n\n2. **Impact of Improved Sequence Level Routing Regularization on Training Stability:**\n   - Does the implementation of improved sequence-level routing regularization enhance the training stability of the model?\n\n3. **Distinct Contributions of MoE Variants in MoEUT:**\n   - Given the implementation of MoE within the MoEUT framework (specifically in the MLP and attention mechanisms), how does each variant individually affect the model's overall performance? \n\n4. **Impact of Excluding MoE on Performance:**\n   - If the MoE design is omitted, allowing layer grouping and “peri-layernorm” to independently influence the model, what is the anticipated impact on performance? Additionally, is the integration of layer grouping and “peri-layernorm” with the proposed MoE architecture necessary for achieving the observed benefits?\n\n5. **Comparative Analysis of MoEUT and SUT:**\n   - In addition to perplexity comparisons, it is suggested that a direct comparison between MoEUT and SUT (Sparse Transformer) should also be conducted across the downstream language modeling tasks presented in Table 1. This would provide a more comprehensive understanding of their relative performances.", "limitations": "The authors adequately discuss the limitations and do not have the potential negative societal impact.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The paper demonstrates the efficacy of incorporating the Mixture of Experts (MoE) architecture into a shared-layer Transformer network, highlighting its potential to improve performance.\n2. By integrating the MoE model with layer grouping and a novel layer normalization approach, the proposed model achieves superior results in language modeling tasks compared to standard Transformers.\n3. Visualization of the results reveals significant specialization across different layers, with each layer showing a distinct preference for particular experts, indicating effective learning and specialization within the network.", "weaknesses": "1. The relationship between the MoE network architecture, the proposed layer grouping strategy, and the novel layer normalization technique appear to be undefined. Clarification on how these components synergize within the model would enhance understanding of their collective impact on performance.\n\n2. While the author suggests that the $\\sigma$-MoE model exhibits instability during training, this assertion is not substantiated with empirical evidence. Providing experimental results or a more detailed analysis to support this claim would strengthen the argument.\n\n3. The current ablation study does not sufficiently demonstrate the effectiveness of the proposed method. Expanding the ablation study to include a broader range of experiments and comparative analyses could offer a more comprehensive evaluation of the individual contributions of the proposed enhancements.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their valuable time reviewing our work. We would like to respond to the concerns raised by the reviewer as follows.\n\n\n> … . Clarification on how these components synergize within the model would enhance understanding of their collective impact on performance…\n\n> Do the layer grouping and the novel layer normalization contribute independently to the performance improvements observed, or are these contributions specifically related to the architecture of the Mixture of Experts (MoE) network?\n\n> If the MoE design is omitted, allowing layer grouping and “peri-layernorm” to independently influence the model, what is the anticipated impact on performance? Additionally, is the integration of layer grouping and “peri-layernorm” with the proposed MoE architecture necessary for achieving the observed benefits?\n\nWe would first like to draw the reviewer’s attention to the fact that we have dedicated ablation studies on the effect of layer grouping in Fig. 6 and the effect of different layer normalizations in Fig. 8 independently of each other for our MoE-based model. We demonstrated that both help.\n\n\nTesting these on the naive UTs without MoE in a systematic/fair way is not easy precisely because of the challenge of scaling the naive UTs (the very problem we address in this work): they are prohibitively slow and use too much memory for any interesting scale. Given that MoE is anyway necessary to scale up UTs, it seemed natural to focus on evaluating these methods for the MoE-based UTs.\n\nAlso intuitively, we have no reason to believe that the contributions of these methods would be very different for UTs because the original motivations of these methods are based on UT’s properties, and they are unrelated to the use of MoE (e.g., peri-layernorm is motivated by the residual norm growth caused by parameter-sharing across layers; see A.2).\n\n\n> The current ablation study does not sufficiently demonstrate the effectiveness of the proposed method. Expanding the ablation study to include a broader range of experiments and comparative analyses could offer a more comprehensive evaluation of the individual contributions of the proposed enhancements.\n\n\nWe believe that the ablation studies presented in the current paper sufficiently cover the most important aspects to justify our design choice within MoEUT. Please see Fig 6 for the effect of routing, Fig 8 for comparing pre/post/peri layernorm, Fig 13 for the effect of d_expert, and Fig 14 for the effect of K. \nAs we explained above, conducting further ablations on the naive UTs without MoE is not reasonable because of their scale inefficiency. If the reviewer still thinks there are any other ablations that are critically missing, we would appreciate it a lot if the reviewer could suggest concrete ideas. Thank you. \n\n> While the author suggests that the 𝜎-MoE model exhibits instability during training, this assertion is not substantiated with empirical evidence. Providing experimental results or a more detailed analysis to support this claim would strengthen the argument.  … Does the implementation of improved sequence-level routing regularization enhance the training stability of the model?\n\nYes, without sequence-level routing regularization, larger models suffer from an expert collapse and they diverge. We found it uninformative to include training loss curves exploding to infinity in the appendix, but if the reviewer thinks that would be interesting for the readers, we can include them in the appendix of the next version of the paper.\n\n> In addition to perplexity comparisons, it is suggested that a direct comparison between MoEUT and SUT (Sparse Transformer) should also be conducted across the downstream language modeling tasks presented in Table 1. This would provide a more comprehensive understanding of their relative performances.\n\n\nThank you for pointing this out. We extended our downstream task evaluations to SUT and “SUT without ACT” (in the meanwhile, we identified that the “ACT” component of SUT, which is a fundamental component of SUT, is detrimental to its performance here). These results show that MoEUT achieves 18-33% increases (8.7-12.2 points) over SUT. Removing ACT from SUT improves its downstream performance as well, but MoEUT remains consistently better, with 7-11% increases (3-5.1 points) on average across our tasks, model sizes, and pretraining datasets.\n\n| Dataset | #params | Model |  PPL | LAMBADA | BLiMP | CBT | HellaSwag | PIQA | ARC-E | Average |\n|----|----|----|----|----|----|----|----|----|----|----|\n| C4 |  44M | MoEUT | 18.30 | 23.2% | 78.2% | 81.1% | 29.2% | 61.3% | 33.5% | 51.1% |\n| | 44M | SUT | 40.50 | 1.2% | 65.3% | 51.1% | 26.4% | 57.8% | 31.9% | 39.0% |\n| | 44M | SUT w.o. ACT | 21.51 | 18.1% | 72.8% | 66.3% | 27.5% | 59.1% | 32.5% | 46.0% |\n| | 244M | MoEUT | 13.24 | 30.6% | 79.7% | 85.3% | 35.7% | 65.2% | 36.4% | 55.5% |\n| | 244M | SUT | 20.05 | 20.5% | 71.0% | 68.5% | 28.2% | 60.1% | 32.7% | 46.8% |\n| | 244M | SUT w.o. ACT | 14.58 | 27.8% | 77.0% | 75.9% | 32.7% | 63.2% | 35.5% | 52.0% |\n| PES2O | 44M | MoEUT | 11.09 | 13.1% | 68.7% | 69.6% | 28.3% | 55.1% | 31.4% | 44.4% |\n| | 44M | SUT | 25.04 | 0.5% | 59.2% | 38.1% | 26.2% | 55.0% | 31.1% | 35.0% |\n| | 44M | SUT w.o. ACT | 12.68 | 11.7% | 66.5% | 53.9% | 28.0% | 56.1% | 31.5% | 41.3% |\n| | 244M | MoEUT | 8.52 | 19.4% | 73.5% | 77.4% | 30.1% | 56.3% | 35.6% | 48.7% |\n| | 244M | SUT | 20.44 | 0.5% | 60.9% | 42.8% | 26.7% | 55.3% | 33.0% | 36.5% |\n| | 244M | SUT w.o. ACT | 9.31 | 16.8% | 71.9% | 64.8% | 28.8% | 57.3% | 34.8% | 45.7% |\n\nWe will add this to the updated version of our paper.\n\nWe believe our response above resolves all the concerns that the reviewer has raised. If the reviewer finds our response useful, please consider increasing the score. Thank you very much.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "OsS1EJqLiN", "reviewer": "Reviewer_gEb4", "summary": "Motivated by the superior generalization performance of Universal Transformers (UT) demonstrated in other works, this paper addresses the compute efficiency problem of this architecture. While UT decreases the parameter count drastically by sharing parameters across layers, vanilla UT underperforms dense transformers on typical NLP tasks. It is unclear how to efficiently scale the parameter count of the UT, as scaling a single layer to compensate for the loss of parameters would result in computational costs beyond those of standard transformers due to increased layer width.\n\nFirst, this paper explores Mixture-of-Experts (MoE) techniques to efficiently scale the parameters of the UT. Due to sparse expert activation, MoE can scale up the parameter count without drastically increasing the compute requirements. Since simply scaling the parameter count with MoE layers does not yield the expected performance gains, the authors propose two additional innovations: Layer grouping—where parameters are shared within groups of layers instead of across all layers—and Peri-layernorm—where layer normalization is applied only before linear projections that are followed by sigmoid or softmax activation functions.\n\nAt a high level, my understanding is that the proposed MoEUT architecture essentially interpolates between a dense transformer without shared layers and a UT with completely shared layers. MoEUT can learn to reuse the same experts for each layer (or group), which corresponds to vanilla UT, or it can learn to use completely different, non-overlapping experts at each layer, resembling a standard transformer.", "questions": "- ll. 28 - 29: the fact that UT shines in systematic generalization baselines does not make it a more \"general\" architecture, does it? In contrary, UT imposes the requirement of parameter sharing across layers, which can be see as an additional inductive bias, which could be potentially learned by the standard transformer from the data.\n- Chapter 3: I am a bit confused by the fact that σ-MoE underperforms the dense transformer baseline in parameter-matched setting. (Fig.4a), which seem to be in contradiction with the results from the σ-MoE paper. Why is it the case?\n- l. 172: my understanding from the related literature is that relative positional encoding is essential to enable systematic generalization in UT. Why do authors decide to use RoPE positional encodings here?\n- l. 219: \"We note that SUTs have not been evaluated previously on any language modeling tasks.\" -- the original SUT was evaluated on e.g. English-German translation task, Is it not a language modelling task?\n- ll. 319 - 322: Authors talk about importance of comparing parameter-matched MoE and dense models. IT would be useful if they could explain why parameter-matched setting is important (as compared to compute matched)\n- ll. 287 - 299: Does \"column\" refer to individual token representations throughout the layers in transformer?", "limitations": "Authors adequately address the limitations in chapter 5.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Overall it is an interesting paper, the proposed method is (conceptually) simple, sound and well motivated. I especially like that authors also included zero-shot downstream task evaluation in the empirical section.\n\nOriginality: the paper builds on existing knowledge in a meaningful way. It uses existing MoE technique to efficiently increase capacity of UT, which is not a new idea as it had already been proposed in the SUT paper (as aknowledged and discussed by the authors). It adds two additional ideas that seem to have a positive effect on the performance: layer grouping and \"peri-layernorm\". Yet, as discussed in chapter 5, also layer-grouping has been investigated in prior works.\n\nQuality & Clarity: overall, the claims of the paper are supported by evidences. The paper is adequately written, though chapters 2.1 and 2.2 might be dense for readers unfamiliar with the two architectures the authors build upon. The arguments are mostly clear with some exceptions (see Questions).\n\nSignificance: good. the paper addresses relevant problem and makes a moderate contribution to the field. This work has the potential to influence subsequent studies on scaling UT architectures.", "weaknesses": "At a high level, my main concern is the following: given that MoEUT effectively interpolates between a standard transformer and vanilla UT, do the authors think that MoEUT would still keep the advantage of vanilla UT when it comes to systematic generalization?\nBeyond that, I have some additional questions/doubts that I list in the questions part.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for the valuable review and for positive comments on the methodology of the paper. Please find our responses as follows:\n\n> .. given that MoEUT effectively interpolates between a standard transformer and vanilla UT, do the authors think that MoEUT would still keep the advantage of vanilla UT when it comes to systematic generalization? \n\nWe can confirm that: we evaluated MoEUT with G=2 on the publicly available dataset of [1] for learning random knowledge graphs. Similarly to Fig. 14 in the appendix of [1], our model learns ~70% OOD generalization. In contrast, the standard Transformers completely fail.\n\n> the fact that UT shines in systematic generalization baselines does not make it a more \"general\" architecture, does it? In contrary, UT imposes the requirement of parameter sharing across layers, which can be see as an additional inductive bias, which could be potentially learned by the standard transformer from the data.\n\nWe would first like to clarify that UTs are in principle strictly more general than *parameter equivalent* standard transformers: they could assign Nparams/Nlayers of their weights to each layer deterministically and simulate a standard Transformer. In this sense, UTs are not at all a “restricted version” of the standard Transformers.\n\nThat said, the reviewer’s remark regarding the standard Transformer and learning layer-sharing from the data is a valid point. In practice with finite amounts of data, however, learning such behavior purely from data is very inefficient, if not impossible. A good illustration of this is compositional tasks. Transformers typically learn to solve compositional tasks by allocating each step/function to each layer. In realistic settings, not all compositional combinations are present in the dataset, thus, certain functions are only learned in certain layers; causing failures to generalize on unseen compositions.\n\nIn fact, [1] rigorously analyzes this in synthetic knowledge graphs, unveiling that transformers use early layers for resolving the first hop, and later layers to resolve the second. However, the knowledge about rarely-seen compositions will be only available in the early layers, thus they can’t be composed with others. [2] shows similar issues on LLama3-70b on real world problems.\n\n> I am a bit confused by the fact that σ-MoE underperforms the dense transformer baseline in parameter-matched setting. (Fig.4a), which seem to be in contradiction with the results from the σ-MoE paper. Why is it the case?\n\nPlease note that our baselines are much stronger than those reported in the 𝜎-MoE paper: on C4, we achieve a perplexity of 13.4 using 244M parameters vs. 17.79 reported by the 𝜎-MoE paper using 266M.\n\nThis difference comes from two modifications. First, the 𝜎-MoE paper follows the experimental protocol of Transformer XL: we used their official 𝜎-MoE codebase, but improved their baseline by using RoPE and no XL cache. Second, they use dropout in the FFN layers of their baseline, and ‘expert dropout’ in the 𝜎-MoE. Here we disabled all dropouts in all our models as we use sub-epoch training. This resulted in perplexity improvements with a higher gain for the baseline than for 𝜎-MoE. \n\n> Why do authors decide to use RoPE positional encodings here?\n\nRoPE is a form of relative positional encoding and is used by most modern LLMs, like LLama. Given its popularity, we considered it the best choice.\n\n\n> \"We note that SUTs have not been evaluated previously on any language modeling tasks.\" -- the original SUT was evaluated on e.g. English-German translation task, Is it not a language modelling task?\n\nWe can soften or remove this claim, since it depends a lot on how one approaches the problem. Translation tasks (like many NLP tasks) can be formulated as an LM task, but in practice, the translation benchmarks (used in the original SUT paper) and those for evaluating LLMs (used in ours) are very different in important ways. The former is an isolated task where people explicitly train models on a dedicated translation dataset (even a specific encoder-decoder architecture is often used instead of generic decoder-only auto-regressive models). This makes translation benchmarks somewhat easier in the sense that very small, specialized models perform pretty well (e.g., the Transformer baseline used in the SUT paper only has 65M parameters; while their biggest SUT also only has 110M params).\n\n> IT would be useful if they could explain why parameter-matched setting is important (as compared to compute matched)\n\nThe parameter-matched setting is crucial to evaluate the model’s *expressiveness* in the LLM tasks where the number of parameters has a high impact on the model performance. We consider this setting to be particularly important to evaluate the true expressiveness of MoEs compared to their dense counterparts.\n\nWhile the compute-matched setup has values when considering certain practical settings, it gives an “unfair” advantage to MoEs in terms of comparison, as we can easily add extra parameters to an MoE without significantly increasing compute requirements. Here we wanted to show that our MoEUT is capable, even without considering such an advantage, by evaluating its pure expressiveness in the more challenging parameter-matched setting.\n\n> Does \"column\" refer to individual token representations throughout the layers in transformer?\n\nYes, it does. We thank the reviewer for pointing out this ambiguity and will improve the clarity in the final version.\n\nWe believe our response above resolves all the concerns that the reviewer has raised. If the reviewer finds our response useful, please consider increasing the score. Thank you very much.\n\n[1] Wang et al: Grokked Transformers are Implicit Reasoners: A Mechanistic Journey to the Edge of Generalization\n\n[2] Biran et al: Hopping Too Late: Exploring the Limitations of Large Language Models on Multi-Hop Queries", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper introduces mixture-of-experts (MoE) in the universal Transformer (UT) architecture with parameter sharing across layers.  UTs have a limitation on parameter-compute ratio.  This parameter count problem makes UTs disadvantageous when dealing with parameter-dominant tasks such as modern language modeling.   Specifically, the authors introduce MoE in both attention and feed-forward layers of UT.  In addition, the authors propose two techniques in the so-called MoEUT. One is layer grouping and the other is to relocate the layer normalization. Experiments shows MoEUT outperforms standard Transformers slightly while incurring less compute and memory costs.  In the rebuttal and follow-up discussions, the authors have addressed most of the concerns by the reviewers with additional experimental results and elaborated explanation. Overall this is an interesting paper and all reviewers recommend acceptance.  That being said, the performance of the proposed MoEUT trained on large scale data and models is still yet to to seen and further technical improvements may need to be investigated.
Decision: Accept (poster)