Research map / Architectures and efficiency / Mixture of experts
MoEUT: Mixture-of-Experts Universal Transformers
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes MoEUT, a mixture-of-experts shared-layer Universal Transformer combining sigma-MoE feedforward, SwitchHead attention, layer grouping, and a peri-layernorm scheme, which matches or slightly beats parameter-matched dense Transformers on language modeling at lower compute.
Problem
Universal Transformers share parameters across layers, which cuts parameter count for a given width and makes naive scaling compute-prohibitive, so they have not been competitive on parameter-dominated tasks like language modeling.
Approach
Make the shared layer mixture-of-experts in both feedforward and attention blocks, recurrently stack groups of non-shared MoE layers (layer grouping), and use a peri-layernorm scheme that applies layernorm only before sigmoid/softmax-gated linear layers and leaves the residual path free of norms.
Key results
- On C4, the 244M MoEUT reaches perplexity 13.24 vs 13.40 for the parameter-matched dense Transformer, and the 1040M MoEUT reaches 10.90 vs 11.15.
- The 244M MoEUT trains 1.7x slower than the dense Transformer (772 vs 443 ms/batch) with a similar memory footprint (9.0G vs 9.2G), and is much faster and lighter than SUT (1344 ms, 23.4G).
- Zero-shot average across six tasks is marginally higher for MoEUT in most settings (e.g., 51.1% vs 49.4% at 44M on C4), with differences generally small.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| replace the dense Transformer with parameter-matched MoEUT | Transformer language models · C4, 44M–1040M | PPL: 18.97 -> 18.30; 13.40 -> 13.24; 11.15 -> 10.90 | Helped |
| replace the dense Transformer with parameter-matched MoEUT | Transformer language models · peS2o | PPL: 11.46 -> 11.09; 8.55 -> 8.52 | Helped |
| replace the dense Transformer with parameter-matched MoEUT | Transformer language models · SlimPajama, 1040M | PPL: 9.56 -> 9.36 | Helped |
| remove ACT (ablation) | SUT · C4, 44M | PPL: 40.50 -> 21.51 | Helped |
| remove ACT (ablation) | SUT · C4, 244M | PPL: 20.05 -> 14.58 | Helped |
| use recurrent layer grouping with group size G=2 instead of G=1 or G=18 (ablation) | MoEUT, 244M · C4 | PPL | Helped |
| stack grouped layers as ABAB instead of AABB (ablation) | MoEUT, 244M, G=2 · C4 | PPL | Helped |
| use peri-layernorm instead of pre-layernorm or post-layernorm (ablation) | MoEUT · C4 | PPL: 11.29 vs 11.32 | Helped |
| use peri-layernorm instead of post-layernorm (ablation) | MoEUT, 719M · C4 | PPL: 11.29 perplexity points compared to 11.32 | Helped |
| use smaller MLP expert sizes while keeping compute fixed (ablation) | MoEUT, 244M · C4 | PPL | Helped |
| increase the number of active MLP experts (ablation) | MoEUT, 244M · C4 | PPL | Helped |
| replace MoEUT with a non-shared Transformer | MoEUT · C4 | ms/batch: 772 vs 443 | Helped |
Surprising findings
- The SUT baseline performs much worse, and most of its poor performance was attributed to its ACT mechanism; removing ACT improved it but it still underperformed MoEUT and the dense baseline.
- Downstream zero-shot differences between MoEUT and the dense baseline are marginal, and at the largest scales the peri-norm advantage over post-norm nearly vanishes (11.29 vs 11.32 PPL at 719M).
Limitations
- The current Triton MoE kernel is suboptimal, making training 1.5-2x slower than dense Transformers with FlashAttention.
- Experiments use a modest training budget (10^5 steps); behavior at massive scale is untested.
- Evaluated only up to about 1B parameters, and the downstream gains are marginal.
Experimental setup
Datasets: C4 SlimPajama peS2o The Stack
Benchmarks: C4 held-out perplexity SlimPajama peS2o The Stack LAMBADA BLiMP Children's Book Test HellaSwag PIQA ARC-E
Metrics: perplexity accuracy MAC operations wall-clock time per batch memory usage
Baselines: dense pre-layernorm Transformer sigma-MoE (non-shared) Sparse Universal Transformer (SUT) Naive UT MoEUT with post-layernorm MoEUT with pre-layernorm
Builds on: sigma-MoE SwitchHead Universal Transformer pre-layernorm post-layernorm
Closest papers by meaning
- Mixture of Tokens: Continuous MoE through Cross-Example Aggregation (NeurIPS 2024 · similarity 0.89)
- SwitchHead: Accelerating Transformers with Mixture-of-Experts Attention (NeurIPS 2024 · similarity 0.88)
- Pruning and Distilling Mixture-of-Experts into Dense Language Models (NeurIPS 2026 · similarity 0.88)
- Monet: Mixture of Monosemantic Experts for Transformers (ICLR 2025 · similarity 0.87)
- From Sparse to Soft Mixtures of Experts (ICLR 2024 · similarity 0.87)
- UMoE: Unifying Attention and FFN with Shared Experts (NeurIPS 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.