UMoE: Unificatie van aandacht en FFN met gedeelde experts

UMoE: Unifying Attention and FFN with Shared Experts

May 12, 2025

Auteurs: Yuanhang Yang, Chaozheng Wang, Jing Li

cs.AI

Samenvatting

Sparse Mixture of Experts (MoE)-architecturen zijn naar voren gekomen als een veelbelovende aanpak voor het schalen van Transformer-modellen. Terwijl initiële werken MoE voornamelijk integreerden in feedforward netwerk (FFN)-lagen, hebben recente studies het MoE-paradigma uitgebreid naar aandachtslagen om de modelprestaties te verbeteren. Bestaande op aandacht gebaseerde MoE-lagen vereisen echter gespecialiseerde implementaties en vertonen suboptimale prestaties in vergelijking met hun op FFN gebaseerde tegenhangers. In dit artikel streven we ernaar de MoE-ontwerpen in aandacht- en FFN-lagen te verenigen door een nieuwe herformulering van het aandachtmechanisme te introduceren, waarbij een onderliggende FFN-achtige structuur binnen aandachtmodules wordt onthuld. Onze voorgestelde architectuur, UMoE, bereikt superieure prestaties via op aandacht gebaseerde MoE-lagen, terwijl efficiënte parametersharing tussen FFN- en aandachtcomponenten mogelijk wordt gemaakt.

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.