UMoE: Unificación de la Atención y la FFN con Expertos Compartidos

UMoE: Unifying Attention and FFN with Shared Experts

May 12, 2025

Autores: Yuanhang Yang, Chaozheng Wang, Jing Li

cs.AI

Resumen

Las arquitecturas de Mezcla Escasa de Expertos (MoE, por sus siglas en inglés) han surgido como un enfoque prometedor para escalar los modelos Transformer. Si bien los trabajos iniciales incorporaron principalmente MoE en las capas de redes neuronales feed-forward (FFN), estudios recientes han explorado la extensión del paradigma MoE a las capas de atención para mejorar el rendimiento del modelo. Sin embargo, las capas MoE basadas en atención existentes requieren implementaciones especializadas y demuestran un rendimiento subóptimo en comparación con sus contrapartes basadas en FFN. En este artículo, buscamos unificar los diseños MoE en las capas de atención y FFN mediante la introducción de una nueva reformulación del mecanismo de atención, revelando una estructura subyacente similar a FFN dentro de los módulos de atención. Nuestra arquitectura propuesta, UMoE, logra un rendimiento superior a través de capas MoE basadas en atención, al mismo tiempo que permite un uso eficiente de parámetros compartidos entre los componentes de FFN y atención.

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.