UMoE: Unificazione dell'Attenzione e della FFN con Esperti Condivisi

UMoE: Unifying Attention and FFN with Shared Experts

May 12, 2025

Autori: Yuanhang Yang, Chaozheng Wang, Jing Li

cs.AI

Abstract

Le architetture Sparse Mixture of Experts (MoE) si sono affermate come un approccio promettente per scalare i modelli Transformer. Mentre i lavori iniziali integravano principalmente MoE nei livelli di rete feed-forward (FFN), studi recenti hanno esplorato l'estensione del paradigma MoE ai livelli di attenzione per migliorare le prestazioni del modello. Tuttavia, gli attuali livelli MoE basati sull'attenzione richiedono implementazioni specializzate e dimostrano prestazioni subottimali rispetto alle loro controparti basate su FFN. In questo articolo, miriamo a unificare i design MoE nei livelli di attenzione e FFN introducendo una nuova riformulazione del meccanismo di attenzione, rivelando una struttura sottostante simile a FFN all'interno dei moduli di attenzione. La nostra architettura proposta, UMoE, raggiunge prestazioni superiori attraverso livelli MoE basati sull'attenzione, consentendo al contempo una condivisione efficiente dei parametri tra i componenti FFN e di attenzione.

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.