MoEUT: Mixture-of-Experts Universal Transformers

Authors: róbert  csordás
kazuki  irie
jürgen  schmidhuber
christopher  potts
christopher d manning
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 16:49:12

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.

Main paper: 
 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 † Work started at IDSIA. 1 Our code is public: https://github.com/robertcsordas/moeut 2 Dehghani et al. [10] also augment UTs with an additional adaptive computation time (ACT, [17, 18] ) mechanism. However, the benefits of UTs we discuss here are purely due to layer-sharing, which, in consequence, is the focus of this work. Our models could also optionally be augmented with ACT but this is out of scope here. 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. 
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.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 dmodel×dexpert and W e 2 ∈ R dexpert×dmodel , where e ∈ {1, . . . , N E } is the index of the expert. At each token position t, given layer input x t ∈ R dmodel , the MoE feedforward layer computes a score for each expert, yielding a vector s ∈ R N E computed as:EQUATIONwhere W S ∈ R dmodel×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 top-K elements in s t ∈ R N E to produce the layer output y t ∈ R dmodel as follows:EQUATIONEQUATIONwhere 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:L = N E e=1 p[e] log p[e], p = 1 T T t=1 softmax(x t W S ) ∈ R N E (4)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.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 dmodel , 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 dmodel×dhead, where h ∈ {1, . . . , H} is the head index. The expert selection for the values is computed as follows:EQUATIONEQUATIONwhere W h SV ∈ R dmodel×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 matrixW h SO ∈ R dmodel×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 dmodel is calculated as follows:EQUATIONEQUATIONEQUATIONwhere W h,e V ∈ R dmodel×dhead and W h,e O ∈ R dhead×dmodel 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] withK h t = (k h 1 , . . . , k h t ), V h t = (v h 1 , . . . , v h t ) ∈ R T ×dhead 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.)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.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] . 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.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. 
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 2d 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 = 2d 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 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, 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 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.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? 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). 
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.So far, all the results have been cumulative statistics in different input sequences and positions.We might wonder about the selection behavior for the individual Transformer columns. 3Is 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. 
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.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. 
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. 
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? 
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). 
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. 
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. 
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. 
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. 
 
 Appendix: 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.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. 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. 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 . 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. 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. 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.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).By representing the Transformer's activations in a 2D grid, with token positions on the x-axis and depth on the y-axis, "columns" correspond to all hidden activations across depth given a token position. 


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

Reviews: review: paper_summary: paper_summary: The paper suggests to use Sparse MoEs together with Universal Transformers to overcome the parameter-count limitation that the latter have when parameters are shared over consecutive layers. In particular, the work suggests to use $\sigma$-MoEs (that use sigmoid activation function in the router, rather than the more popular softmax activation), and use fine-grained experts (selecting many small experts over a large pool, rather than one or two large ones over a smaller pool). The goal is to make Universal Transformers competitive on language modeling tasks, where they haven’t excelled in the past (presumably due to the smaller parameter count, compared to dense models of the same time complexity).

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

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


-------
Edit 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".
questions: questions: - In figure 10 and 11, do the expert indices correspond to the experts in the attention or in the feed-forward part of the block? Do the trends differ across the two?
limitations: limitations: No negative societal impact particular to this work, in my opinion.
paper_summary: paper_summary: This paper focuses on the problem of inefficient parameter-computation ratio in Universal Transformers (UT). UT shares parameters across layers but reduces the parameter count significantly. One naive approach is to scaling up the layer size, however, cannot be easily achieved due to the prohibitive computational resource requirements. MoEUT, the proposed approach, exploits MoE architecture in both attention and feed forward layers to address this issue. To achieve the similar performance as the standard Transformers, the authors first introduce layer grouping, which allows non-shared weights among the layers within a group. Second, they remove the layernorm in the main data path in standard Transformer to avoid gradient flow and add layernormaround the residual connections. Experiments shows MoEUT outperforms standard Transformers slightly while incurring less compute and memory costs.
strength_weakness: strength_weakness: 1. The claim of using significantly less compute and memory should be further supported by evaluation numbers. Though Table 4 have presented detailed training hardware information of the experiments reported in the paper, it is very hard to determine the quantitative computation and memory costs. A general number like GPU hours, GPU duty cycles, DRAM utilization or costs in terms of dollars would be a much better indicator rather than existing presentation.
2. 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.
questions: questions: It seems the number of experts is very large as a default setting in the evaluation. What if the total number of experts is just, let's say, 8 or 16 and k is 1 or 2, similar like the sparsely-activated MoE, how the performance would change?
limitations: limitations: Limitations are addressed in the paper and there is no negative societal impact of this work.
paper_summary: paper_summary: This paper introduces a novel application of the mixture of experts (MoE) architecture within both the MLP and attention modules of the Universal Transformer network. The integration of MoE is further complemented by an innovative sequence-level routing regularization technique, which the authors argue enhances training stability. Additionally, the paper proposes the incorporation of layer grouping and a new layer normalization scheme, aimed at boosting model performance. Compared with baselines, the proposed MoEUT can achieve better perplexity with the same number of parameters and obtain competitive performance in various downstream tasks.
strength_weakness: strength_weakness: 1. The relationship between the MoE network architecture, the proposed layer grouping strategy, and the novel layer normalization technique appear to be undefined. Clarification on how these components synergize within the model would enhance understanding of their collective impact on performance.

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

3. 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.
questions: questions: 1. **Layer Grouping and Novel Layer Norm Performance Contribution:**
   - 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?

2. **Impact of Improved Sequence Level Routing Regularization on Training Stability:**
   - Does the implementation of improved sequence-level routing regularization enhance the training stability of the model?

3. **Distinct Contributions of MoE Variants in MoEUT:**
   - 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? 

4. **Impact of Excluding MoE on Performance:**
   - 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?

5. **Comparative Analysis of MoEUT and SUT:**
   - In addition to perplexity comparisons, it is suggested that a direct comparison between MoEUT and SUT (Sparse Transformer) should also be conducted across the downstream language modeling tasks presented in Table 1. This would provide a more comprehensive understanding of their relative performances.
limitations: limitations: The authors adequately discuss the limitations and do not have the potential negative societal impact.
paper_summary: paper_summary: Motivated by the superior generalization performance of Universal Transformers (UT) demonstrated in other works, this paper addresses the compute efficiency problem of this architecture. While UT decreases the parameter count drastically by sharing parameters across layers, vanilla UT underperforms dense transformers on typical NLP tasks. It is unclear how to efficiently scale the parameter count of the UT, as scaling a single layer to compensate for the loss of parameters would result in computational costs beyond those of standard transformers due to increased layer width.

First, 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.

At 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.
strength_weakness: strength_weakness: At a high level, my main concern is the following: given that MoEUT effectively interpolates between a standard transformer and vanilla UT, do the authors think that MoEUT would still keep the advantage of vanilla UT when it comes to systematic generalization?
Beyond that, I have some additional questions/doubts that I list in the questions part.
questions: questions: - ll. 28 - 29: the fact that UT shines in systematic generalization baselines does not make it a more "general" architecture, does it? In contrary, UT imposes the requirement of parameter sharing across layers, which can be see as an additional inductive bias, which could be potentially learned by the standard transformer from the data.
- 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?
- 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?
- 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?
- 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)
- ll. 287 - 299: Does "column" refer to individual token representations throughout the layers in transformer?
limitations: limitations: Authors adequately address the limitations in chapter 5.

Comments: title: please respond to rebuttal
Official Comment by Reviewer ZbxN
comment: Thank you very much for the increased score! We are glad to hear that the reviewer found our response useful! Thank you again for your valuable feedback!
I thank the authors very much for their detailed response to my comments, and the other reviewers' comments as well.

In particular, I appreciate very much the fact that the authors made the effort to fairly reproduce the baselines, and contacted the authors to make sure that they were correctly represented.

Thus, given that I feel that my concerns have mostly been addressed, I will increase the score and recommend the acceptance of the paper.
Thank you very much for your response! We are glad to hear that the reviewer found our response useful!
We are very thankful for the increased score! We are glad to hear that the reviewer found our response useful!

We agree with the reviewer that analyzing PeriLN in more detail would be a very useful additional experiment for our paper. First, we would like to emphasize that one direction of such comparison is already part of the paper: we analyze the effect of different layer norm variants on MoEUT in Figure 8 of the paper. The reverse direction is missing: We do not have experiments with PeriLN on standard (naive) UTs. We tried to run such experiments upon the suggestion of the reviewer. We would like to emphasize that this is very resource intensive: we need at least 8 A6000-level GPUs to make the experiment fit in the memory even for our tiny, 44M naive UT model. Unfortunately, we found that naive UTs experience residual blowups and become unstable during training. 
Such blowup does not happen in MoEs, possibly because the number of simultaneously active channels/attention heads is limited to a much smaller number than in naive UTs, where all of them can be active at the same time, resulting in significantly greater residual norm growth. We added an additional regularization on the norm of the residual, which delayed the blowup significantly, but it still happened at around 20% of the training completed. Until that point its behavior was as we predicted: naive UT with PeriLN outperformed naive UT without PeriLN (using PreLN instead). We are happy to finalize these experiments, search for a stable solution, and add them to the final version of the paper, but since naive UT is a completely novel setting that we have not examined so far, unfortunately, the remaining time is not enough for us to do it during the rebuttal.

We would like to thank the reviewer again for his efforts to improve the quality of our paper and for his thoughtful suggestions!
Dear reviewers,

The author-reviewer discussion period will end soon. Please make sure you read the authors' rebuttal and respond to it. If you have additional questions after reading the rebuttal please discuss with the authors. For those who have done so, thank you!

AC
Thank you very much for addressing my concerns. They are really helpful. I will maintain my current rating.
I extend my gratitude to the authors for their detailed responses. To fully grasp the efficacy and interplay of the proposed components, I anticipate an exploration of the performance for MoEUT without PeriLN and vice versa. Such results would provide a clearer understanding of these elements.

Moreover, I am keen to understand how MoEUT and PeriLN complement each other, suggesting that their integration is not merely coincidental but inherently synergistic.

The authors have addressed my concerns to some extent, therefore I would like to increase my score. Should the authors further elaborate on these aspects, specifically by demonstrating the independent and combined effects of MoEUT and PeriLN, I am inclined to further improve their score.
Thank you very much for your response! We are glad to hear that the reviewer found our response useful!
I would like to thank the authors for their comprehensive replies. I appreciate that the authors tested MoEUT on systematic generalization tasks (learning random graphs), making the method even more convincing.