UMoE: Unifying Attention and FFN with Shared Experts

TL;DR

UMoE addresses Transformer scalability by unifying MoE designs for attention and FFN through a reformulation that reveals an FFN-like structure within attention. It introduces shared expert components across both modules and a pre-mixing attention variant with low-rank query projections, achieving superior performance to FFN-based MoEs while maintaining the same parameter count. Experiments on FineWeb-Edu and Wikitext-103 show improved perplexity and zero-shot accuracy, with attention-based MoEs matching or surpassing FFN-based MoEs and clear evidence of expert specialization. The work highlights practical efficiency via MAC-based analysis and points to future directions in more efficient token mixing and universal transformer architectures.

Abstract

Sparse Mixture of Experts (MoE) architectures have emerged as a promising approach for scaling Transformer models. While initial works primarily incorporated MoE into feed-forward network (FFN) layers, recent studies have explored extending the MoE paradigm to attention layers to enhance model performance. However, existing attention-based MoE layers require specialized implementations and demonstrate suboptimal performance compared to their FFN-based counterparts. In this paper, we aim to unify MoE designs in attention and FFN layers by introducing a novel reformulation of the attention mechanism, that reveals an underlying FFN-like structure within attention modules. Our proposed architecture, UMoE, achieves superior performance through attention-based MoE layers while enabling efficient parameter sharing between FFN and attention components.

Figures (10)

• Figure 1: Illustration of a UMoE layer, which incorporates MoE into both FFN and attention modules with shared experts. The primary distinction between attention-MoE and FFN-MoE lies in an additional token mixing operation.
• Figure 2: Two formulations of the multi-head attention mechanism. (a) Vanilla attention interleaves mixing operations with value and output projections. (b) Pre-mixing attention performs token mixing prior to projections.
• Figure 3: Implementation details of a UMoE layer. The input consists of a sequence X containing n token hidden states and x representing the final hidden state. For simplicity, this implementation focuses on computing the output for the last token.
• Figure 4: Post-Mixing Attention.
• Figure 5: Best valid PPL (top) and training loss (bottom) on Wikitext.