Unión de Expertos: Adaptación del Enrutamiento Jerárquico al Transformador Descompuesto de Forma Equivalente

Union of Experts: Adapting Hierarchical Routing to Equivalently Decomposed Transformer

March 4, 2025

Autores: Yujiao Yang, Jing Lian, Linhui Li

cs.AI

Resumen

Mixture-of-Experts (MoE) mejora el rendimiento del modelo mientras mantiene la eficiencia computacional, lo que lo hace adecuado para aplicaciones a gran escala. Sin embargo, en el paradigma actual de MoE, cada experto funciona de manera individual, careciendo así de interacciones de alta calidad entre expertos. Además, no se han extendido eficazmente a los bloques de atención, lo que limita mejoras adicionales en la eficiencia. Para abordar estos problemas, proponemos Union-of-Experts (UoE), que descompone el transformador en un grupo equivalente de expertos y luego implementa enrutamiento dinámico sobre los datos de entrada y los expertos. Nuestro enfoque avanza el diseño de MoE con tres innovaciones clave: (1) Realizamos una descomposición equivalente de expertos tanto en bloques MLP como en bloques de atención basada en la partición de matrices en paralelismo tensorial. (2) Desarrollamos dos paradigmas de enrutamiento: selección de datos por parches y selección de expertos, para aplicar el enrutamiento en diferentes niveles. (3) Diseñamos la arquitectura del modelo UoE, incluyendo Atención Multi-Cabezal Selectiva (SMHA) y Union-of-MLP-Experts (UoME). (4) Implementamos en paralelo las operaciones de enrutamiento y cálculo de UoE, y optimizamos la eficiencia basándonos en el análisis del procesamiento hardware. Los experimentos demuestran que el modelo equipado con UoE supera a Full Attention, a los MoE de última generación y a los transformadores eficientes en varias tareas de los dominios de imágenes y lenguaje natural. Los códigos fuente están disponibles en https://github.com/YujiaoYang-work/UoE.

English

Mixture-of-Experts (MoE) enhances model performance while maintaining computational efficiency, making it well-suited for large-scale applications. However, expert in exist MoE paradigm works as an individual, thereby lacking high-quality expert interactions. Moreover, they have not been effectively extended to attention block, which constrains further efficiency improvements. To tackle these issues, we propose Union-of-Experts (UoE), which decomposes transformer into an equitant group of experts, and then implement dynamic routing on input data and experts. Our approach advances MoE design with three key innovations: (1) We conducted equitant expert decomposition on both MLP blocks and attention blocks based on matrix partition in tensor parallelism. (2) We developed two routing paradigms: patch wise data selection and expert selection, to apply routing across different levels. (3) We design the architecture of UoE model, including Selective Multi-Head Attention (SMHA) and Union-of-MLP-Experts (UoME). (4) We develop parallel implementation of UoE's routing and computation operation, and optimize efficiency based on the hardware processing analysis. The experiments demonstrate that the model employed with UoE surpass Full Attention, state-of-art MoEs and efficient transformers in several tasks across image and natural language domains. The source codes are available at https://github.com/YujiaoYang-work/UoE.