MoEUT: Mixture-of-Experts Universal Transformers

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.

Strengths

Overall it is an interesting paper, the proposed method is (conceptually) simple, sound and well motivated. I especially like that authors also included zero-shot downstream task evaluation in the empirical section. Originality: 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. Quality & 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). Significance: good. the paper addresses relevant problem and makes a moderate contribution to the field. This work has the potential to influence subsequent studies on scaling UT architectures.

Weaknesses

At a high level, my main concern is the following: given that MoEUT effectively interpolates between a standard transformer and vanilla UT, do the authors think that MoEUT would still keep the advantage of vanilla UT when it comes to systematic generalization? Beyond that, I have some additional questions/doubts that I list in the questions part.

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?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Authors adequately address the limitations in chapter 5.

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.

Strengths

1. The paper demonstrates the efficacy of incorporating the Mixture of Experts (MoE) architecture into a shared-layer Transformer network, highlighting its potential to improve performance. 2. 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. 3. Visualization of the results reveals significant specialization across different layers, with each layer showing a distinct preference for particular experts, indicating effective learning and specialization within the network.

Weaknesses

1. The relationship between the MoE network architecture, the proposed layer grouping strategy, and the novel layer normalization technique appear to be undefined. Clarification on how these components synergize within the model would enhance understanding of their collective impact on performance. 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

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.

Rating

6

Confidence

4

Soundness

2

Presentation

3

Contribution

2

Limitations

The authors adequately discuss the limitations and do not have the potential negative societal impact.

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.

Strengths

1. The idea of exploiting MoE to address the parameter-compute ration in Universal Transformers is interesting. MoE itself incorporates ideas of expert sharing among similar tokens. It is an intuitively reasonable combination. 2. 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. 3. The evaluation and analysis are comprehensive with deep dive to the design components.

Weaknesses

1. The claim of using significantly less compute and memory should be further supported by evaluation numbers. Though Table 4 have presented detailed training hardware information of the experiments reported in the paper, it is very hard to determine the quantitative computation and memory costs. A general number like GPU hours, GPU duty cycles, DRAM utilization or costs in terms of dollars would be a much better indicator rather than existing presentation. 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

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?

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Limitations are addressed in the paper and there is no negative societal impact of this work.

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.

Strengths

- The presentation of the work is excellent. From the motivation, to the explanation of the proposed method, related works, and the description of the experiments. - The main experiments show that the goal of making UTs competitive with standard dense transformers in language modelling tasks was achieved. - Many additional ablation experiments were conducted to try to explain how experts are selected across different layers, and tokens. - An effort was put on implementing reasonable baselines such as $\sigma$-MoEs and Sparse UTs. - The paper clearly states current limitations (e.g. not the most efficient implementation, which results in experiments that are 2x slower than the dense counterpart).

Weaknesses

- The proposed method matches dense transformers on language modeling, but it’s barely better. This begs the question: why using MoEUTs then, rather than the simpler baseline? The answer could be better MAC efficiency at the same memory cost, but this is not clearly represented in Table 1 (and as the authors point out, it’s actually slower due to implementation limitations). - 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

- 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?

Rating

7

Confidence

5

Soundness

3

Presentation

4

Contribution

3

Limitations

No negative societal impact particular to this work, in my opinion.

Official Comment by Reviewer ZbxN

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.

please respond to rebuttal

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