UMoE: Unifying Attention and FFN with Shared Experts

Jing Li (The Hong Kong Polytechnic University) · Yuanhang Yang (Institute of Science Tokyo) · Chaozheng Wang (Department of Computer Science and Engineering, The Chinese University of Hong Kong)

attention layers attention mechanismattention-based moeefficient architecturefeed-forward networkffn-based counterpartsffn-like structuremodel performance model scalability parameter sharingreformulationsparse mixture of expertsspecialized implementationstransformer modelsumoe architecture

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.