Research map / Architectures and efficiency / Mixture of experts
UMoE: Unifying Attention and FFN with Shared Experts
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 UMoE, which reformulates multi-head attention as pre-mixing token aggregation followed by FFN-like experts so attention and FFN layers can share one MoE expert set, improving perplexity and zero-shot accuracy over FFN-MoE, MoA and SwitchHead at 134M and 1.1B scales.
Problem
Attention-based MoE layers underperform FFN-based MoE and require specialized implementations; the paper asks whether attention can be reformulated to share the FFN expert design without losing expressiveness.
Approach
Reformulate attention as pre-mixing (token mixing first, then grouped value and output projections), which is an FFN-like two-matrix expert applied to mixed inputs. Routing these experts in both attention and FFN layers lets one shared expert pool serve both modules, with low-rank expert-specific query projections.
Key results
- At base scale (134M dense), UMoE reaches 20.44 FineWeb-Edu PPL versus 21.19 for fine-grained FFN-MoE and 22.28 for MoA, with 540M total params.
- At 1.1B dense scale, UMoE achieves 15.95 FineWeb-Edu PPL vs 16.09 for FFN-MoE, and the highest zero-shot average (47.58 vs 47.12 for FFN-MoE).
- UMoE-Att (attention-only MoE) outperforms MoA and SwitchHead at base scale (20.81 vs 22.28 and 22.91 PPL) without sharing parameters with FFN.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use UMoE instead of Fine-grained FFN-MoE | Base language model · FineWeb-Edu | perplexity: 21.19 -> 20.44 | Helped |
| Use UMoE instead of Fine-grained FFN-MoE | Base language model · Wikitext-103 | perplexity: 27.94 -> 26.67 | Helped |
| Use UMoE instead of Fine-grained FFN-MoE | Large language model · FineWeb-Edu | perplexity: 16.09 -> 15.95 | Helped |
| Use UMoE instead of MoA | Large language model · Wikitext-103 | perplexity: 25.14 -> 25.44 | Hurt |
| Use UMoE instead of Fine-grained FFN-MoE | Base model · Zero-shot average | accuracy: $39.55$ -> $40.06$ | Helped |
| Use UMoE instead of Fine-grained FFN-MoE | Large model · Zero-shot average | accuracy: $47.12$ -> $47.58$ | Helped |
| Shift expert allocation from FFN layers toward attention layers (ablation) | UMoE · FineWeb-Edu (20B tokens) | perplexity: 22.82 -> 21.75 | Helped |
| Add activation functions between expert matrix multiplications (ablation) | UMoE · FineWeb-Edu (20B tokens) | perplexity: 24.43 -> 22.82 | Helped |
| Add activation functions between expert matrix multiplications (ablation) | UMoE-Att · FineWeb-Edu (20B tokens) | perplexity: 23.99 -> 23.37 | Helped |
| Use separate fixed experts and routers across attention and FFN layers (ablation) | UMoE · FineWeb-Edu (20B tokens) | perplexity: $23.11$ -> $22.82$ | Helped |
| Replace pre-mixing attention with post-mixing attention (ablation) | UMoE attention MoE · Wikitext-103 and FineWeb-Edu | language-modeling performance | Hurt |
Surprising findings
- Removing activation functions from experts only modestly degrades results (PPL 24.43 vs 22.82), and UMoE remains trainable without them.
- Allocating all 20 activated experts to attention yields the best perplexity (21.75) in the expert allocation ablation, suggesting attention is more expressive than FFN.
Limitations
- Modest extra computation from pre-mixing token mixing in small models.
- Experiments do not cover mathematics or code, and reasoning ability is not evaluated.
- No error bars or statistical significance tests due to compute constraints.
Experimental setup
Datasets: FineWeb-Edu 100B Wikitext-103
Benchmarks: FineWeb-Edu validation perplexity Wikitext-103 perplexity HellaSwag PIQA ARC-E ARC-C RACE Lambada MMLU Winogrande
Metrics: perplexity zero-shot accuracy MACs pre-filling latency training throughput
Baselines: Dense Transformer Fine-grained FFN-MoE MoA SwitchHead
Builds on: Mixture of Experts Mixture-of-Attention (MoA) SwitchHead DeepSeekMoE Switch Transformer load balancing loss
Closest papers by meaning
- SwitchHead: Accelerating Transformers with Mixture-of-Experts Attention (NeurIPS 2024 · similarity 0.90)
- A unified framework for establishing the universal approximation of transformer-type architectures (NeurIPS 2025 · similarity 0.87)
- MoEUT: Mixture-of-Experts Universal Transformers (NeurIPS 2024 · similarity 0.87)
- Mixture-of-Top-$k$ Attention: Efficient Attention as Scalable Fast Weights (NeurIPS 2026 · similarity 0.87)
- Rethinking Attention: Exploring Shallow Feed-Forward Neural Networks as an Alternative to Attention Layers in Transformers (Student Abstract) (AAAI 2024 · similarity 0.87)
- From Sparse to Soft Mixtures of Experts (ICLR 2024 · similarity 0.86)
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.