CompeteSMoE -- Entrenamiento de Mezcla de Expertos con Garantía Estadística mediante Competencia

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

Resumen

La mezcla dispersa de expertos (SMoE, por sus siglas en inglés) ofrece una solución atractiva para escalar la complejidad del modelo más allá de simplemente aumentar la profundidad o el ancho de la red. Sin embargo, argumentamos que el entrenamiento efectivo de SMoE sigue siendo un desafío debido al proceso de enrutamiento subóptimo, donde los expertos que realizan cálculos no contribuyen directamente al proceso de enrutamiento. En este trabajo, proponemos la competencia, un mecanismo novedoso para dirigir tokens a los expertos con la respuesta neuronal más alta. Teóricamente, demostramos que el mecanismo de competencia tiene una mejor eficiencia muestral que el enrutamiento tradicional basado en softmax. Además, desarrollamos CompeteSMoE, un algoritmo simple pero efectivo para entrenar modelos de lenguaje grandes mediante la implementación de un enrutador que aprende la política de competencia, logrando así un alto rendimiento con un bajo costo de entrenamiento. Nuestras extensas evaluaciones empíricas en tareas de ajuste visual por instrucción y preentrenamiento de lenguaje demuestran la eficacia, robustez y escalabilidad de CompeteSMoE en comparación con las estrategias SMoE más avanzadas. Hemos puesto a disposición la implementación en: https://github.com/Fsoft-AIC/CompeteSMoE. Este trabajo es una versión mejorada del estudio previo en 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