União de Especialistas: Adaptando o Roteamento Hierárquico para Transformadores Equivalentemente Decompostos

Union of Experts: Adapting Hierarchical Routing to Equivalently Decomposed Transformer

March 4, 2025

Autores: Yujiao Yang, Jing Lian, Linhui Li

cs.AI

Resumo

O Mixture-of-Experts (MoE) melhora o desempenho do modelo enquanto mantém a eficiência computacional, tornando-o bem adequado para aplicações em grande escala. No entanto, no paradigma MoE existente, cada especialista funciona de forma individual, carecendo de interações de alta qualidade entre os especialistas. Além disso, eles não foram efetivamente estendidos para blocos de atenção, o que limita melhorias adicionais de eficiência. Para resolver esses problemas, propomos o Union-of-Experts (UoE), que decompõe o transformer em um grupo equivalente de especialistas e, em seguida, implementa o roteamento dinâmico nos dados de entrada e nos especialistas. Nossa abordagem avança o design do MoE com três inovações principais: (1) Realizamos a decomposição equivalente de especialistas tanto em blocos MLP quanto em blocos de atenção com base na partição de matrizes no paralelismo tensorial. (2) Desenvolvemos dois paradigmas de roteamento: seleção de dados por patch e seleção de especialistas, para aplicar o roteamento em diferentes níveis. (3) Projetamos a arquitetura do modelo UoE, incluindo a Selective Multi-Head Attention (SMHA) e o Union-of-MLP-Experts (UoME). (4) Desenvolvemos a implementação paralela das operações de roteamento e computação do UoE e otimizamos a eficiência com base na análise de processamento de hardware. Os experimentos demonstram que o modelo empregado com UoE supera a Full Attention, os MoEs state-of-the-art e os transformers eficientes em várias tarefas nos domínios de imagem e linguagem natural. Os códigos-fonte estão disponíveis em 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.