MoEUT: Mixture-of-Experts Universal Transformers

We would like to thank the reviewer for the valuable review and for positive comments on the methodology of the paper. Please find our responses as follows:

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

We can confirm that: we evaluated MoEUT with G=2 on the publicly available dataset of [1] for learning random knowledge graphs. Similarly to Fig. 14 in the appendix of [1], our model learns ~70% OOD generalization. In contrast, the standard Transformers completely fail.

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.

We would first like to clarify that UTs are in principle strictly more general than parameter equivalent standard transformers: they could assign Nparams/Nlayers of their weights to each layer deterministically and simulate a standard Transformer. In this sense, UTs are not at all a “restricted version” of the standard Transformers.

That said, the reviewer’s remark regarding the standard Transformer and learning layer-sharing from the data is a valid point. In practice with finite amounts of data, however, learning such behavior purely from data is very inefficient, if not impossible. A good illustration of this is compositional tasks. Transformers typically learn to solve compositional tasks by allocating each step/function to each layer. In realistic settings, not all compositional combinations are present in the dataset, thus, certain functions are only learned in certain layers; causing failures to generalize on unseen compositions.

In fact, [1] rigorously analyzes this in synthetic knowledge graphs, unveiling that transformers use early layers for resolving the first hop, and later layers to resolve the second. However, the knowledge about rarely-seen compositions will be only available in the early layers, thus they can’t be composed with others. [2] shows similar issues on LLama3-70b on real world problems.

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?

Please note that our baselines are much stronger than those reported in the 𝜎-MoE paper: on C4, we achieve a perplexity of 13.4 using 244M parameters vs. 17.79 reported by the 𝜎-MoE paper using 266M.

This difference comes from two modifications. First, the 𝜎-MoE paper follows the experimental protocol of Transformer XL: we used their official 𝜎-MoE codebase, but improved their baseline by using RoPE and no XL cache. Second, they use dropout in the FFN layers of their baseline, and ‘expert dropout’ in the 𝜎-MoE. Here we disabled all dropouts in all our models as we use sub-epoch training. This resulted in perplexity improvements with a higher gain for the baseline than for 𝜎-MoE.

Why do authors decide to use RoPE positional encodings here?

RoPE is a form of relative positional encoding and is used by most modern LLMs, like LLama. Given its popularity, we considered it the best choice.

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

We can soften or remove this claim, since it depends a lot on how one approaches the problem. Translation tasks (like many NLP tasks) can be formulated as an LM task, but in practice, the translation benchmarks (used in the original SUT paper) and those for evaluating LLMs (used in ours) are very different in important ways. The former is an isolated task where people explicitly train models on a dedicated translation dataset (even a specific encoder-decoder architecture is often used instead of generic decoder-only auto-regressive models). This makes translation benchmarks somewhat easier in the sense that very small, specialized models perform pretty well (e.g., the Transformer baseline used in the SUT paper only has 65M parameters; while their biggest SUT also only has 110M params).

IT would be useful if they could explain why parameter-matched setting is important (as compared to compute matched)

The parameter-matched setting is crucial to evaluate the model’s expressiveness in the LLM tasks where the number of parameters has a high impact on the model performance. We consider this setting to be particularly important to evaluate the true expressiveness of MoEs compared to their dense counterparts.

While the compute-matched setup has values when considering certain practical settings, it gives an “unfair” advantage to MoEs in terms of comparison, as we can easily add extra parameters to an MoE without significantly increasing compute requirements. Here we wanted to show that our MoEUT is capable, even without considering such an advantage, by evaluating its pure expressiveness in the more challenging parameter-matched setting.

Does "column" refer to individual token representations throughout the layers in transformer?

Yes, it does. We thank the reviewer for pointing out this ambiguity and will improve the clarity in the final version.

We believe our response above resolves all the concerns that the reviewer has raised. If the reviewer finds our response useful, please consider increasing the score. Thank you very much.

[1] Wang et al: Grokked Transformers are Implicit Reasoners: A Mechanistic Journey to the Edge of Generalization

[2] Biran et al: Hopping Too Late: Exploring the Limitations of Large Language Models on Multi-Hop Queries

We would like to thank the reviewer for their valuable time reviewing our work. We would like to respond to the concerns raised by the reviewer as follows.

… . Clarification on how these components synergize within the model would enhance understanding of their collective impact on performance…

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?

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?

We would first like to draw the reviewer’s attention to the fact that we have dedicated ablation studies on the effect of layer grouping in Fig. 6 and the effect of different layer normalizations in Fig. 8 independently of each other for our MoE-based model. We demonstrated that both help.

Testing these on the naive UTs without MoE in a systematic/fair way is not easy precisely because of the challenge of scaling the naive UTs (the very problem we address in this work): they are prohibitively slow and use too much memory for any interesting scale. Given that MoE is anyway necessary to scale up UTs, it seemed natural to focus on evaluating these methods for the MoE-based UTs.

Also intuitively, we have no reason to believe that the contributions of these methods would be very different for UTs because the original motivations of these methods are based on UT’s properties, and they are unrelated to the use of MoE (e.g., peri-layernorm is motivated by the residual norm growth caused by parameter-sharing across layers; see A.2).

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.

We believe that the ablation studies presented in the current paper sufficiently cover the most important aspects to justify our design choice within MoEUT. Please see Fig 6 for the effect of routing, Fig 8 for comparing pre/post/peri layernorm, Fig 13 for the effect of d_expert, and Fig 14 for the effect of K. As we explained above, conducting further ablations on the naive UTs without MoE is not reasonable because of their scale inefficiency. If the reviewer still thinks there are any other ablations that are critically missing, we would appreciate it a lot if the reviewer could suggest concrete ideas. Thank you.

While the author suggests that the 𝜎-MoE model exhibits instability during training, this assertion is not substantiated with empirical evidence. Providing experimental results or a more detailed analysis to support this claim would strengthen the argument. … Does the implementation of improved sequence-level routing regularization enhance the training stability of the model?

Yes, without sequence-level routing regularization, larger models suffer from an expert collapse and they diverge. We found it uninformative to include training loss curves exploding to infinity in the appendix, but if the reviewer thinks that would be interesting for the readers, we can include them in the appendix of the next version of the paper.

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.

Thank you for pointing this out. We extended our downstream task evaluations to SUT and “SUT without ACT” (in the meanwhile, we identified that the “ACT” component of SUT, which is a fundamental component of SUT, is detrimental to its performance here). These results show that MoEUT achieves 18-33% increases (8.7-12.2 points) over SUT. Removing ACT from SUT improves its downstream performance as well, but MoEUT remains consistently better, with 7-11% increases (3-5.1 points) on average across our tasks, model sizes, and pretraining datasets.

We will add this to the updated version of our paper.

We believe our response above resolves all the concerns that the reviewer has raised. If the reviewer finds our response useful, please consider increasing the score. Thank you very much.

We would like to thank the reviewer for the valuable review and for many positive comments on the methodology of our paper. Please find our responses as follows:

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.

We will update Table 4 in the final version if the paper is accepted, to include the total GPU-hours (N_GPU * duration in the current table) for better readability. Comparing memory this way is nontrivial, as it does not scale linearly with the number of GPUs, as we are using simple data-parallel training. However, here are some direct comparisons on identical hardware for the 244M parameter models that allow directly comparing them on identical hardware, showing that MoEUT is much faster and uses less memory than the alternative UT variants:

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.

I may not fully understand the difference here and why full batch balancing loss would cause explosion while sequence-level won't.

We cannot provide a mathematical proof here, but we do have one hypothetical but intuitive explanation. If the balancing is done at the batch level, the experts can specialize in different sequence types. Since the batches are sampled randomly, the “variance” of utilization rate between the active experts used in different forward passes might vary highly. Some experts may specialize in rare sequences, thus they will be rarely trained, resulting in a “slow expert collapse”. When they are reactivated at a much later stage of training, they may cause sudden loss explosion as their training “stage” lags behind that of other more frequently used experts. On the other hand, if the balancing is done on a sequence level, the model is encouraged to use various experts in each sequence, eliminating some of this variance. This still allows some expert specialization, because the balancing is relatively weak.

Using a significantly larger batch size would probably also stabilize the training, but it is prohibitively expensive with our resource budget. It seems like a useful next step for better-resourced teams, and so we will emphasize this in our next version.

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?

This is an excellent question. In our experiments, we found K=1 performed significantly worse than K > 1. In order to match the MACs of the dense model, the number of experts should be at least K*N_layers, otherwise the activated experts would be “wider” than the FFN of the baseline. Furthermore, to keep the number of parameters constant, d_expert should be significantly increased.

We study the effect of increasing expert size (d_expert) in Fig. 13 in the appendix, while keeping K constant. Increasing the expert size is detrimental. Decreasing K would have an additional negative effect. We analyze this independently of K in Fig 14.

If the reviewer thinks additional experiments could add value to the paper, we can run an experiment with some of our MoEUT methods in a max-d_expert configuration and K=2 to see the effects of pushing these parameters to the extreme.

These questions are extremely valuable for us and will enable us to improve our paper. Thank you very much.

We would like to thank the reviewer for their insightful review and for the positive comments on the clarity and methodology of the paper. Please find our responses as follows:

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?

Our goal is to pave the way towards a general model that can be used for LM but is also good at systematic generalization, eventually arriving at foundation models that are more trustworthy and data efficient. In this paper, we address a single step toward that goal: we present a way to overcome the long-standing efficiency limitations of UTs (which are known to generalize better than standard transformers), which have prevented their scaling to sizes required for LLMs.

it’s actually slower due to implementation limitations

It is true that our current proof-of-concept kernel is slower than a standard transformer despite requiring fewer flops, but we believe this can be mostly mitigated by better kernels. Even with the current kernels, our method is significantly faster than a naive UT or the next best UT baseline, SUT, and for most of our experiments naive UTs are not viable to run because of their memory usage and speed (see time/memory usage measurements below). Rather than replacing Transformers, we are advocating for more research on Universal Transformers with our method for larger scales.

Note additionally that we investigated these models in a parameter-matched setting because we were interested in comparing their expressiveness. The number of parameters in our MoEUT can be increased very cheaply (as with other MoEs), and it is possible to achieve significantly better perplexity for marginally higher compute, potentially justifying the wall-clock time slowdown even with the current kernels. (We use a parameter-matched setup instead of this because this would give an unfair advantage to our MoEUTs compared to the baselines).

Speed/memory usage of different models, showing that MoEUT is much faster compared to naive UT and SUT and uses less memory:

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.

The perplexity reported for Sparse UTs is alarmingly high

We confirmed via personal communication through a colleague that SUT does not work well for LM. We used the official code of the authors, with minimal modification required to adapt it to our codebase. Additionally, we conducted some ablations on the SUT models, and we found that the main cause of the bad performance is the ACT used by their model (which was presented as a fundamental building block of SUT). By disabling ACT, the gap between MoEUT and SUT is much smaller, while MoEUT remains consistently better:

A similar thing happens with 𝜎-MoEs, which performs worse than the dense baseline when matching the parameter count.

Please note that our baselines are much stronger than those reported in the 𝜎-MoE paper: on C4, we achieve a perplexity of 13.4 using 244M parameters vs. 17.79 reported by the 𝜎-MoE paper using 266M.

This difference comes from two modifications. First, the 𝜎-MoE paper follows the experimental protocol of Transformer XL: we used their official 𝜎-MoE codebase, but improved their baseline by using RoPE and no XL cache. Second, they use dropout in the FFN layers of their baseline, and ‘expert dropout’ in the 𝜎-MoE. Here we disabled all dropouts in all our models as we use sub-epoch training. This resulted in perplexity improvements with a higher gain for the baseline than for 𝜎-MoE.

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?

Those expert indices correspond to the feed-forward part. We will clarify this in the final version. Thank you for pointing this out!

We believe our response above resolves all the concerns that the reviewer has raised. If the reviewer finds our response useful, please consider increasing the score. Thank you very much.