UMoE: 共有エキスパートによるアテンションとFFNの統合

UMoE: Unifying Attention and FFN with Shared Experts

May 12, 2025

著者: Yuanhang Yang, Chaozheng Wang, Jing Li

cs.AI

要旨

スパースなMixture of Experts（MoE）アーキテクチャは、Transformerモデルのスケーリングにおいて有望なアプローチとして登場した。初期の研究では主にフィードフォワードネットワーク（FFN）層にMoEを組み込んでいたが、最近の研究ではモデルの性能向上を目指し、MoEパラダイムをアテンション層に拡張することが検討されている。しかし、既存のアテンションベースのMoE層は専用の実装を必要とし、FFNベースのMoE層と比較して最適な性能を発揮していない。本論文では、アテンションメカニズムの新たな再定式化を導入し、アテンションモジュール内に潜在するFFNに類似した構造を明らかにすることで、アテンション層とFFN層におけるMoE設計を統一することを目指す。提案するアーキテクチャであるUMoEは、アテンションベースのMoE層を通じて優れた性能を達成しつつ、FFNとアテンションコンポーネント間の効率的なパラメータ共有を可能にする。

English

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 the MoE designs in attention and FFN layers by introducing a novel reformulation of the attention mechanism, revealing 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.