UMoE: Unifying Attention and FFN with Shared Experts

TL;DR

UMoE introduces a unified approach to sparse Mixture of Experts in Transformers by reformulating attention mechanisms to share parameters with FFN layers, leading to improved performance and efficiency.

Contribution

The paper presents UMoE, a novel architecture that unifies attention and FFN MoE layers through a reformulated attention mechanism enabling shared experts.

Findings

UMoE outperforms traditional MoE models in accuracy.

Shared experts reduce model complexity and improve efficiency.

Unified design simplifies implementation of MoE in Transformers.

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.