Title: MoEUT: Mixture-of-Experts Universal Transformers
Venue: NeurIPS 2024
Accepted: yes
Keywords: MoE

mixture of experts

universal transformers

language modeling
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

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

1 Introduction

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

One important candidate is the Universal Transformer (UT, [10]). The core characteristic of UTs is recurrence in depth via sharing parameters across layers . This reintroduces the expressive power of recurrence provided by recurrent neural networks (RNNs, [11, 12, 13]). Layer sharing allows UTs to outperform regular Transformers on compositional problems such as logical inference tasks, while also yielding improvements on small-scale language modeling and translation tasks. In particular, UTs have been shown to have better compositional generalization properties [14, 15] by being able to decompose structured problems without supervision and generalize to longer sequences [16]. 2 These empirical findings confirm that UTs are more general architectures with superior

generalization properties compared to standard Transformers, in principle. However, UTs suffer from a fundamental problem of parameter-compute ratio : sharing the parameters among L layers of an L -layer Transformer-while keeping the same model dimensionalities-results in a model with L times 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 L times wider) usually yields a very big layer whose computational requirements in terms of compute and memory are prohibitive in practice [19, 20]. In sum, despite their potential, UTs are much less compute-efficient than standard Transformers, and thus, they are not popular for parameter-dominated tasks such as modern language modeling. Indeed, we are not aware of any previous work that has succeeded in developing compute-efficient UT models that yield competitive performance compared to standard Transformers on such tasks.

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

MoEUTs allow us to build parameter- and resource-efficient UT language models outperforming standard Transformers with less compute and memory requirements on all scales on which we can afford to test (up to 1B parameters). We demonstrate 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 we address the parameter-compute ratio problem by using mixture-of-experts. While there are many recent works on MoE methods for Transformer language models (e.g., [24, 25, 26, 27, 28]), making them competitive against their dense counterparts in parameter-equal comparisons is known to be challenging [28]. Here we leverage recent advances in MoE methods for both the feedforward network block (FFN, or simply MLP layer or feedforward layer; Sec. 2.1) and the self-attention layer (Sec. 2.2) together with two novel methods that take into account the specific properties of shared-layer models, namely: layer grouping (Sec. 2.3) and signal propagation (Sec. 2.4), which, taken together, are crucial for achieving effective shared-layer MoE Transformers.

2.1 MoE Feedforward Blocks

To parameterize the feedforward blocks of our shared-layer Transformers by an MoE, we use σ -MoE [28] with a few modifications. σ -MoE divides the feedforward block into N E slices, called experts . Each expert has two sets of weights, W e 1 ∈ R d model × d expert and W e 2 ∈ R d expert × d model , where e ∈ { 1 , . . . , N E } is the index of the expert. At each token position t , given layer input x t ∈ R d model , the MoE feedforward layer computes a score for each expert, yielding a vector s ∈ R N E computed as:

where W S ∈ R d model × N E is a trainable weight matrix, and σ ( x ) = 1 1+ e -x is the element-wise sigmoid function. The MoE layer only selects K experts (out of N E ) corresponding to the topK elements in s t ∈ R N E to produce the layer output y t ∈ R d model as follows:

where s t [ e ] ∈ R is the e -th element of vector s t ∈ R N E . Our preliminary experiments revealed that the original regularization of σ -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 x t , t ∈ { 1 , . . . , T } we compute the balancing loss L as:

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

2.2 MoE Self-Attention Layers

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

where W h SV ∈ R d model × N A is the selection weight for the value and K A is the number of simultaneously active experts per head, set to K A = 2 in all of our experiments. The selection for the values and outputs are independent. The selection of the output is computed analogously using a different weight matrix W h SO ∈ R d model × N A : s h O,t = σ ( x t W h SO ) ∈ R N A and E h O ( x t ) = arg topk( s h O,t , K A ) ⊂ { 1 , . . . , N A } . Then the output y ∈ R d model is calculated as follows:

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

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

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

detrimental. Second, the total number of attention heads is kept relatively low, which might not be sufficient for a large model. Increasing their number is similarly prohibitively expensive.

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

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

2.4 Novel LayerNorm Scheme for Improved Signal Propagation in Universal Transformers

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

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

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

0

200

400

600

800

1000

Number of model parameters (M)

12

14

16

18

Perplexity

Baseline

MoEUT

σ

-MoE

0

2

4

6

Total training MACs

×

10 18

12

14

16

18

Perplexity

Baseline

MoEUT

σ

-MoE

(a) Scaling in the number of parameters.

(b) Scaling in the number of MACs used for training.

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

3 Main Experimental Results

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

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

244M

728M

Number of parameters

2.50

2.55

2.60

2.65

2.70

2.75

2.80

Perplexity

Baseline

MoEUT

AABB

G=1

G=2

G=3

G=6

G=9

G=18

Variant

13.2

13.4

13.6

13.8

14.0

14.2

Perplexity

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 on a held-out subset of C4. The plot shows that our MoEUT model slightly outperforms dense models with the same number of parameters (Fig. 4a), and the gap tends to grow with scale. Additionally, we compare to the non-shared σ -MoE model [28]. This σ -MoE baseline has the same shape of feedforward layers ( d model, K , d expert) to our layer-shared MoEUT, but uses no attention experts to keep the proportion of the attention weights as close MoEUT as possible, and it also uses our peri-layernorm scheme (Sec. 2.4). We add this baseline as the model that is as close to our shared-layer model as possible. This model performs significantly worse than MoEUT, demonstrating the clear advantage of the shared layers. Additionally, Fig. 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.

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

Zero-shot performance on downstream tasks. Here we evaluate the zero-shot performance of our models on six different downstream tasks: LAMBADA [44], BLiMP [45], Children's Book Test (CBT) [46], HellaSwag [47], PIQA [48], and ARC-E [49]. For LAMBADA, we use the detokenized version from OpenAI, and we evaluate the top-1 accuracy of the last word (it can span multiple tokens; here we use greedy decoding). For CBT and BLiMP, we measure the accuracy for each task and report the average of the tasks' accuracies. The results are shown in Tab. 1. We observe that our models and the baselines typically perform very similarly. MoEUT often outperforms the baseline, but the differences are marginal in all cases. This confirms that our models are indeed very capable compared to standard language models. We confirm this on peS2o and SlimPajama as 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 on 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-layernorm and does not use layer grouping. Unlike ours, Adaptive Computation Time (ACT) is used in the layer dimension.

We took the original code released by Tan et al. [38] and ported it to our training pipeline for a fair comparison. As for MoEUT, we roughly match the model's dimensionalities and number of active channels to our dense baselines. We ran a hyperparameter optimization for the regularization losses,

and we found that a minimal regularization is necessary for stabilizing the training. However, larger regularization tends to hurt performance significantly. All other hyperparameters are set based on 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 to our MoEUT and the parameter-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.

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

Evaluating layernorm schemes. Here we evaluate our 'peri-layernorm' scheme (Sec. 2.4). Fig. 8 shows the results. The proposed layernorm scheme consistently performs the best. The gap is more significant for the small models, while for the bigger ones the gains diminish (for the 719M-parameter

C4, 44M

C4, 244M

peS2o, 44M

peS2o, 244M

8

10

20

30

40

Perplexity

Baseline

MoEUT

SUT w.o. ACT

SUT

44M

244M

13

14

15

16

17

18

19

Perplexity

PostLN

PreLN

PeriLN

model, the gap between peri-norm and post-norm is marginal: 11.29 perplexity points compared to 11.32). At the same time, we also observe that the gap between peri-norm and post-norm increases with the number of training steps, leaving open the possibility of higher gains if the models are trained longer. All our MoEUT models in other experiments make use of this peri-norm scheme.

4 Analysis

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

Expert (re)use across layers. We first focus on whether nontrivial reuse occurs between different layers. Note that MoE-based shared-layer models could, in theory, assign different experts to different layers to 'emulate' regular non-shared models. If this were the case, the model would resemble a regular Transformer instead of a Universal one. To confirm that our model is more versatile than that, we analyze whether certain experts 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 to visualize the distribution of layers that each expert prefers. To better visualize the structure, we reorder experts by a heuristic which we call 'layer position score' defined as the average of the layer indices weighted by the number of expert activations in that layer. Fig. 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 is a wide range of layers where the expert is activated. Experts seem to be active in a continuous sequence of layers. This can be seen in the wide, vertically lined structure. We can conclude that MoEUT is capable of specializing in a specific layer if necessary and sharing weights between them when advantageous.

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

Expert selection dynamics in individual columns/positions. So far, all the results have been cumulative statistics in different input sequences and positions. We might

0

100

200

300

Expert ID

1

3

5

7

9

11

13

15

17

Layer

0.00

0.25

0.50

0.75

0

200

400

600

800

1000

Token

20

40

60

#Experts Used

Layer 1

Layer 5

Layer 9

Layer 13

Layer 17

wonder about the selection behavior for the individual Transformer columns. 3 Is expert selection mostly constant throughout the layers for individual columns of MoEUT?

To answer this question, we calculate the pairwise intersection-overunion 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 a non-negligible overlap between the selected experts in subsequent layers; however, it is far from complete overlap. This indicates that experts usually change dynamically in a single column, performing different functionality in different layers.

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

5 Discussion and Limitations

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

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

1

5

9

13

17

Layer

1

5

9

13

17

Layer

0.25

0.50

0.75

1.00

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 sharedlayer Transformers, also showing that layer grouping outperforms vanilla Universal Transformers. However, they consider models with large group sizes ( G = 6 for 12 layers) and few recurrent steps (2). We find that models with smaller G with more steps perform better. Sometimes, layer grouping is also used to up-scaling pretrained models [55].

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

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

6 Conclusion

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

References

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.

Jürgen Schmidhuber. Learning to control fast-weight memories: An alternative to recurrent nets. Neural Computation , 4(1):131-139, 1992.

Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.

Tom B Brown et al. Language models are few-shot learners. In Proc. Advances in Neural Information Processing Systems (NeurIPS) , Virtual only, December 2020.

OpenAI. ChatGPT. https://openai.com/blog/chatgpt , 2022.

OpenAI. GPT-4 technical report. Preprint arXiv:2303.08774 , 2023.

Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurélien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. LLaMA: Open and efficient foundation language models. Preprint arXiv:2302.13971 , 2023.

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.

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.

Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal Transformers. In Int. Conf. on Learning Representations (ICLR) , New Orleans, LA, USA, May 2019.

Jeffrey L. Elman. Finding structure in time. Cognitive Science , 14(2):179-211, 1990.

Michael I Jordan. Attractor dynamics and parallelism in a connectionist sequential machine. In Proc. Conf. of the Cognitive Science Society , pages 531-546. Amherst, MA, USA, August 1986.

Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation , pages 1735-1780, 1997.

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.

Róbert Csordás, Kazuki Irie, and Jürgen Schmidhuber. The devil is in the detail: Simple tricks improve systematic generalization of Transformers. In Proc. Conf. on Empirical Methods in Natural Language Processing (EMNLP) , Punta Cana, Dominican Republic, November 2021.

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.

Jürgen Schmidhuber. Self-delimiting neural networks. Preprint arXiv:1210.0118 , 2012.

Alex Graves. Adaptive computation time for recurrent neural networks. In Int. Conf. on Learning Representations (ICLR) Workshop Track , Vancouver, Canada, April 2016.

Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. Preprint arXiv:2001.08361 , 2020.

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.

Robert A. Jacobs, Michael I. Jordan, Steven J. Nowlan, and Geoffrey E. Hinton. Adaptive mixtures of local experts. Neural Compututaion , 3(1):79-87, 1991.

John B. Hampshire II and Alexander H. Waibel. The meta-pi network: connectionist rapid adaptation for high-performance multi-speaker phoneme recognition. In Proc. IEEE Int. Conf. on Acoustics, Speech and Signal Processing (ICASSP) , pages 165-168, Albuquerque, New Mexico, USA, April 1990.

Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. In Int. Conf. on Learning Representations (ICLR) , Toulon, France, April 2017.

Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. GShard: Scaling giant models with conditional computation and automatic sharding. In Int. Conf. on Learning Representations (ICLR) , Virtual only, May 2021.

William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Preprint arXiv:2101.03961 , 2021.

Aidan Clark, Diego de Las Casas, Aurelia Guy, Arthur Mensch, Michela Paganini, Jordan Hoffmann, Bogdan Damoc, Blake A. Hechtman, Trevor Cai, Sebastian Borgeaud, George van den Driessche, Eliza Rutherford, Tom Hennigan, Matthew Johnson, Katie Millican, Albin Cassirer, Chris Jones, Elena Buchatskaya, David Budden, Laurent Sifre, Simon Osindero, Oriol Vinyals, Jack W. Rae, Erich Elsen, Koray Kavukcuoglu, and Karen Simonyan. Unified scaling laws for routed language models. Preprint arXiv:2202.01169 , 2022.

Xiaofeng Zhang, Yikang Shen, Zeyu Huang, Jie Zhou, Wenge Rong, and Zhang Xiong. Mixture of attention heads: Selecting attention heads per token. In Proc. Conf. on Empirical Methods in Natural Language Processing (EMNLP) , pages 4150-4162, Abu Dhabi, United Arab Emirates, December 2022.

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.

Jakub Krajewski, Jan Ludziejewski, Kamil Adamczewski, Maciej Pióro, Michal Krutul, Szymon Antoniak, Kamil Ciebiera, Krystian Król, Tomasz Odrzygózdz, Piotr Sankowski, Marek Cygan, and Sebastian Jaszczur. Scaling laws for fine-grained mixture of experts. Preprint arXiv:2402.07871 , 2024.

Damai Dai, Chengqi Deng, Chenggang Zhao, R. X. Xu, Huazuo Gao, Deli Chen, Jiashi Li, Wangding Zeng, Xingkai Yu, Y. Wu, Zhenda Xie, Y. K. Li, Panpan Huang, Fuli Luo, Chong Ruan, Zhifang Sui, and Wenfeng Liang. DeepSeekMoE: Towards ultimate expert specialization in mixture-of-experts language models. Preprint arXiv:2401.06066 , 2024.

Róbert Csordás, Piotr Pi˛ ekos, Kazuki Irie, and Jürgen Schmidhuber. SwitchHead: Accelerating transformers with mixture-of-experts attention. In Preprint arXiv:2312.07987 , December 2023.

Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads. Transformer Circuits Thread , 2022. https://transformer-circuits.pub/2022/in-context-learning-and-inductionheads/index.html.

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.

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.

Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. Preprint arXiv:1607.06450 , 2016.

Simon J. Thorpe. Local vs. distributed coding. Intellectica , 8:3-40, 1989.

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.

Shawn Tan, Yikang Shen, Zhenfang Chen, Aaron C. Courville, and Chuang Gan. Sparse universal transformer. In Proc. Conf. on Empirical Methods in Natural Language Processing (EMNLP) , pages 169-179, Singapore, December 2023.

Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of Machine Learning Research (JMLR) , 21:140:1-140:67, 2020.

Daria Soboleva, Faisal Al-Khateeb, Robert Myers, Jacob R Steeves, Joel Hestness, and Nolan Dey. SlimPajama: A 627B token cleaned and deduplicated version of RedPajama, June 2023. URL https://huggingface.co/datasets/cerebras/SlimPajama-627B .

Luca Soldaini and Kyle Lo. peS2o (Pretraining Efficiently on S2ORC) Dataset. Technical report, Allen Institute for AI, 2023. https://github.com/allenai/pes2o .

Denis Kocetkov, Raymond Li, Loubna Ben Allal, Jia Li, Chenghao Mou, Carlos Muñoz Ferrandis, Yacine Jernite, Margaret Mitchell, Sean Hughes, Thomas Wolf, Dzmitry Bahdanau, Leandro von Werra, and Harm de Vries. The Stack: 3 TB of permissively licensed source code. Preprint arXiv:2211.15533 , 2022.

Jianlin Su, Yu Lu, Shengfeng Pan, Bo Wen, and Yunfeng Liu. RoFormer: Enhanced transformer with rotary position embedding. Preprint arXiv:2104.09864 , 2021.

Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Quan Ngoc Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The LAMBADA dataset: Word prediction requiring a broad discourse context. In Proc. Association for Computational Linguistics (ACL) , Berlin, Germany, August 2016.

Alex Warstadt, Alicia Parrish, Haokun Liu, Anhad Mohananey, Wei Peng, Sheng-Fu Wang, and Samuel R. Bowman. BLiMP: The benchmark of linguistic minimal pairs for English. Transactions of the Association for Computational Linguistics (TACL) , 8:377-392, 2020.

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.

Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a machine really finish your sentence? In Proc. Association for Computational Linguistics (ACL) , pages 4791-4800, Florence, Italy, August 2019.

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.

Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try ARC, the AI2 reasoning challenge. Preprint arXiv:1803.05457 , 2018.

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.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional Transformers for language understanding. In Proc. North American Chapter of the Association for Computational Linguistics on Human Language Technologies (NAACL-HLT) , pages 4171-4186, Minneapolis, MN, USA, June 2019.

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.

Shufang Xie, Huishuai Zhang, Junliang Guo, Xu Tan, Jiang Bian, Hany Hassan Awadalla, Arul Menezes, Tao Qin, and Rui Yan. Residual: Transformer with dual residual connections. Preprint arXiv:2304.14802 , 2023.

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.

Dahyun Kim, Chanjun Park, Sanghoon Kim, Wonsung Lee, Wonho Song, Yunsu Kim, Hyeonwoo Kim, Yungi Kim, Hyeonju Lee, Jihoo Kim, Changbae Ahn, Seonghoon Yang, Sukyung Lee, Hyunbyung Park, Gyoungjin Gim, Mikyoung Cha, Hwalsuk Lee, and Sunghun Kim. SOLAR

10.7B: Scaling large language models with simple yet effective depth up-scaling. Preprint arXiv:2312.15166 , 2023.

Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and memory-efficient exact attention with IO-awareness. In Proc. Advances in Neural Information Processing Systems (NeurIPS) , New Orleans, Louisiana, USA, December 2022.

Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. PyTorch: An imperative style, high-performance deep learning library. In Proc. Advances in Neural Information Processing Systems (NeurIPS) , pages 8024-8035, Vancouver, Canada, December 2019.

Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In Int. Conf. on Learning Representations (ICLR) , New Orleans, LA, USA, May 2019.

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.

A Appendix

A.1 Broader Impact Statement

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

A.2 Growing Residual Norm In Standard Transformers

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

0

5

10

15

Layer

0

50

100

Update magnitude

A.3 Zero-Shot Downstream Performance of SUT-variants

In addition to evaluating the perplexity of different SUT variants in Fig. 7, we also show their zero-shot downstream performance in Tab. 2. It can be seen that MoEUT consistently outperforms both SUT and SUT without ACT.

A.4 Hyperparameters

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

A.5 The Effects of d expert and K

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

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

128

256

512

1024

dexpert

13.2

13.3

13.4

13.5

13.6

13.7

Perplexity

8

16

24

K

13.1

13.2

13.3

13.4

13.5

Perplexity

A.6 Additional Analysis

Expert specialization to tokens/layers. Now conversely to the 'per-token expert selection diversity' analysis presented in the main text, we analyze whether the experts activated by a token are layerspecific for that specific token. For this, we 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, we observe that experts tend to be more layer-specific 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 a small fraction of the used experts are present there. This is consistent with the findings for Layer 1 in Fig. 10.

0

1000

2000

3000

4000

5000

Token

1

5

9

13

17

Layer

0.25

0.50

0.75

1.00

A.7 Wall Clock Time and Memory Comparison

To show the real-world resource usage of our models directly, we run a controlled experiment on identical hardware with our 244M parameter model and the corresponding baselines.

We 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. We show the results in Tab. 5.

Even though our MoEUT with the current kernel implementation is slower than the corresponding dense non-shared layer transformer, it is significantly faster and uses much less memory compared to alternative UT variants.

A.8 Compute Requirements

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

Note that the report is generated from Weights and Biases logs. Because of this, the duration might include effects of restarts because of SLURM preemption and effects for occasional slowdowns of the network drive. Additionally, in a few instances the Weights and Biases log buffer overflowed, making it impossible to determine the number of GPUs used for the experiment. In this case, we report '??' for the 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.

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.

It is fine to include aspirational goals as motivation as long as it 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 in 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 on 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 a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a 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 be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [NA]

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

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a 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 or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental 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)?

Answer: [Yes]

Justification: We show all the hyperparameter configurations in Appendix A.4, and we provide the code for our experiments.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to 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 a 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 a 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 it 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 to 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 a 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 submissions 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 to reproduce that algorithm.

(b) If the contribution is primarily a new 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 a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe 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., to 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 instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We provide the code for our experiments. It auto-downloads all the data that it needs.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so '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 to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) 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. If only a 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 as possible in supplemental material (appended to the paper) is 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, hyperparameters, 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 a 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

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [No]

Justification: Our experiments involve large models that are very expensive to train, and we do not have sufficient compute resources to run multiple seeds of them.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence 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.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 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

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We report the type of hardware used for our main experiments in Appendix A.8. We also give a 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 as 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 it 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://neurips.cc/public/EthicsGuidelines ?

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 a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: We consider our paper to 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 or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness 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 to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to 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 as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a 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 a 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 to 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 we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.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 or dataset.

The authors should state which version of the asset is used and, if possible, include a 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, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset'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 to also provide an easy to use, well documented single-file version of our model.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of 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 is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an 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 as details about compensation (if any)?

Answer: [NA]

Justification: We 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 in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in 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 crowdsourcing nor research with human subjects.

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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nMotivated 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.\n\nSTRENGTHS:\nOverall 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.\n\nWEAKNESSES:\nAt 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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nAuthors adequately address the limitations in chapter 5.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. **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.\n\nLIMITATIONS:\nThe authors adequately discuss the limitations and do not have the potential negative societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\nIt 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?\n\nLIMITATIONS:\nLimitations are addressed in the paper and there is no negative societal impact of this work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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).\n\nWEAKNESSES:\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? 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\".\n\nQUESTIONS:\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? Do the trends differ across the two?\n\nLIMITATIONS:\nNo negative societal impact particular to this work, in my opinion.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-18"}