UMoE : Unification de l’attention et des réseaux feed-forward avec des experts partagés

UMoE: Unifying Attention and FFN with Shared Experts

May 12, 2025

Auteurs: Yuanhang Yang, Chaozheng Wang, Jing Li

cs.AI

Résumé

Les architectures de mélange d'experts parcimonieux (MoE) se sont imposées comme une approche prometteuse pour l'extension des modèles Transformer. Alors que les travaux initiaux intégraient principalement le MoE dans les couches de réseaux feed-forward (FFN), des études récentes ont exploré l'extension du paradigme MoE aux couches d'attention afin d'améliorer les performances des modèles. Cependant, les couches MoE basées sur l'attention existantes nécessitent des implémentations spécialisées et présentent des performances sous-optimales par rapport à leurs homologues basées sur les FFN. Dans cet article, nous visons à unifier les conceptions MoE dans les couches d'attention et FFN en introduisant une nouvelle reformulation du mécanisme d'attention, révélant une structure sous-jacente de type FFN au sein des modules d'attention. Notre architecture proposée, UMoE, atteint des performances supérieures grâce à des couches MoE basées sur l'attention tout en permettant un partage efficace des paramètres entre les composants FFN et d'attention.

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.