CompeteSMoE -- Treinamento de Mistura de Especialistas com Garantia Estatística via Competição

CompeteSMoE -- Statistically Guaranteed Mixture of Experts Training via Competition

May 19, 2025

Autores: Nam V. Nguyen, Huy Nguyen, Quang Pham, Van Nguyen, Savitha Ramasamy, Nhat Ho

cs.AI

Resumo

A mistura esparsa de especialistas (SMoE) oferece uma solução atraente para escalar a complexidade do modelo além do método de aumentar a profundidade ou largura da rede. No entanto, argumentamos que o treinamento eficaz de SMoE permanece desafiador devido ao processo de roteamento subótimo, onde os especialistas que realizam a computação não contribuem diretamente para o processo de roteamento. Neste trabalho, propomos a competição, um mecanismo inovador para direcionar tokens aos especialistas com a maior resposta neural. Teoricamente, mostramos que o mecanismo de competição possui uma melhor eficiência amostral em comparação com o roteamento tradicional por softmax. Além disso, desenvolvemos o CompeteSMoE, um algoritmo simples, porém eficaz, para treinar grandes modelos de linguagem ao implantar um roteador para aprender a política de competição, aproveitando assim um forte desempenho com baixo custo de treinamento. Nossas extensas avaliações empíricas em tarefas de ajuste fino visual e pré-treinamento de linguagem demonstram a eficácia, robustez e escalabilidade do CompeteSMoE em comparação com as estratégias SMoE mais avançadas. Disponibilizamos a implementação em: https://github.com/Fsoft-AIC/CompeteSMoE. Este trabalho é uma versão aprimorada do estudo anterior em arXiv:2402.02526.

English

Sparse mixture of experts (SMoE) offers an appealing solution to scale up the model complexity beyond the mean of increasing the network's depth or width. However, we argue that effective SMoE training remains challenging because of the suboptimal routing process where experts that perform computation do not directly contribute to the routing process. In this work, we propose competition, a novel mechanism to route tokens to experts with the highest neural response. Theoretically, we show that the competition mechanism enjoys a better sample efficiency than the traditional softmax routing. Furthermore, we develop CompeteSMoE, a simple yet effective algorithm to train large language models by deploying a router to learn the competition policy, thus enjoying strong performances at a low training overhead. Our extensive empirical evaluations on both the visual instruction tuning and language pre-training tasks demonstrate the efficacy, robustness, and scalability of CompeteSMoE compared to state-of-the-art SMoE strategies. We have made the implementation available at: https://github.com/Fsoft-AIC/CompeteSMoE. This work is an improved version of the previous study at arXiv:2402.02526