UMoE: Unificando Atenção e FFN com Especialistas Compartilhados

UMoE: Unifying Attention and FFN with Shared Experts

May 12, 2025

Autores: Yuanhang Yang, Chaozheng Wang, Jing Li

cs.AI

Resumo

As arquiteturas de Mistura Especializada de Especialistas (MoE, do inglês Mixture of Experts) têm surgido como uma abordagem promissora para escalar modelos Transformer. Enquanto os trabalhos iniciais incorporavam principalmente o MoE nas camadas de redes feed-forward (FFN), estudos recentes têm explorado a extensão do paradigma MoE para as camadas de atenção, visando melhorar o desempenho do modelo. No entanto, as camadas MoE baseadas em atenção existentes exigem implementações especializadas e demonstram desempenho subótimo em comparação com suas contrapartes baseadas em FFN. Neste artigo, buscamos unificar os designs de MoE nas camadas de atenção e FFN, introduzindo uma nova reformulação do mecanismo de atenção, revelando uma estrutura subjacente semelhante à FFN dentro dos módulos de atenção. Nossa arquitetura proposta, UMoE, alcança desempenho superior por meio de camadas MoE baseadas em atenção, ao mesmo tempo em que permite o compartilhamento eficiente de parâmetros entre os componentes de FFN e atenção.

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.