Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

MoEUT: Mixture-of-Experts Universal Transformers

Róbert Csordás, Kazuki Irie, Jürgen Schmidhuber, Christopher Potts, Christopher D. Manning

TL;DR

This work tackles the inefficiency of shared-layer Universal Transformers by introducing MoEUT, a mixture-of-experts UT that decouples parameters from compute. It combines sigma-MoE FFN blocks and SwitchHead MoE attention with two innovations—layer grouping and peri-layernorm—to enable scalable, parameter-efficient UTs. Empirical results on C4, SlimPajama, peS2o, and The Stack show MoEUT matching or slightly surpassing parameter-matched dense Transformers with reduced compute and memory, and strong zero-shot performance. The paper also provides thorough analyses of expert routing and discusses limitations and avenues for massive scaling, suggesting a path for reviving UT research at large scales.

Abstract

Previous work on Universal Transformers (UTs) has demonstrated the importance of parameter sharing across layers. By allowing recurrence in depth, UTs have advantages over standard Transformers in learning compositional generalizations, but layer-sharing comes with a practical limitation of parameter-compute ratio: it drastically reduces the parameter count compared to the non-shared model with the same dimensionality. Naively scaling up the layer size to compensate for the loss of parameters makes its computational resource requirements prohibitive. In practice, no previous work has succeeded in proposing a shared-layer Transformer design that is competitive in parameter count-dominated tasks such as language modeling. Here we propose MoEUT (pronounced "moot"), an effective mixture-of-experts (MoE)-based shared-layer Transformer architecture, which combines several recent advances in MoEs for both feedforward and attention layers of standard Transformers together with novel layer-normalization and grouping schemes that are specific and crucial to UTs. The resulting UT model, for the first time, slightly outperforms standard Transformers on language modeling tasks such as BLiMP and PIQA, while using significantly less compute and memory.

MoEUT: Mixture-of-Experts Universal Transformers

TL;DR

This work tackles the inefficiency of shared-layer Universal Transformers by introducing MoEUT, a mixture-of-experts UT that decouples parameters from compute. It combines sigma-MoE FFN blocks and SwitchHead MoE attention with two innovations—layer grouping and peri-layernorm—to enable scalable, parameter-efficient UTs. Empirical results on C4, SlimPajama, peS2o, and The Stack show MoEUT matching or slightly surpassing parameter-matched dense Transformers with reduced compute and memory, and strong zero-shot performance. The paper also provides thorough analyses of expert routing and discusses limitations and avenues for massive scaling, suggesting a path for reviving UT research at large scales.

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.
Paper Structure (30 sections, 5 equations, 15 figures, 6 tables)

This paper contains 30 sections, 5 equations, 15 figures, 6 tables.

Table of Contents

  1. Introduction
  2. The MoEUT Architecture
  3. MoE Feedforward Blocks
  4. MoE Self-Attention Layers
  5. Layer Grouping: MoE-efficient Layer Sharing & Sub-operations within an Operation
  6. Novel LayerNorm Scheme for Improved Signal Propagation in Universal Transformers
  7. Main Experimental Results
  8. Scaling compared to standard Transformers.
  9. Performance on code generation.
  10. Zero-shot performance on downstream tasks.
  11. Comparing with SUT.
  12. Evaluating layer grouping.
  13. Evaluating layernorm schemes.
  14. Analysis
  15. Discussion and Limitations
  16. ...and 15 more sections

Figures (15)

  • Figure 1: Layer grouping: 8 layers with group size of 2.
  • Figure 2: The residual grows in pre-layernorm transformers.
  • Figure 3: MoEUT block with no layernorms in the residual.
  • Figure 4: Scaling of different models on C4 (with perplexity measured on a held-out subset of C4). (a) MoEUT slightly outperforms parameter-matched models with no layer sharing. The gap grows with scale. (b) Given equal amounts of compute, MoEUT outperforms other models by a large margin.
  • Figure 5: Performance of MoEUT compared to a standard Transformer on The Stack. MoEUT outperforms standard Transformers. The gap grows with scale.
  • ...and 10 more figures