Mezcla de Expertos con Múltiples Cabezas

Multi-Head Mixture-of-Experts

April 23, 2024

Autores: Xun Wu, Shaohan Huang, Wenhui Wang, Furu Wei

cs.AI

Resumen

Las Mezclas Escasas de Expertos (SMoE, por sus siglas en inglés) escalan la capacidad del modelo sin aumentos significativos en los costos de entrenamiento e inferencia, pero presentan los siguientes dos problemas: (1) Baja activación de expertos, donde solo un pequeño subconjunto de expertos se activa para la optimización. (2) Falta de capacidades analíticas detalladas para múltiples conceptos semánticos dentro de tokens individuales. Proponemos la Mezcla de Expertos de Múltiples Cabezas (MH-MoE), que emplea un mecanismo de múltiples cabezas para dividir cada token en múltiples sub-tokens. Estos sub-tokens son asignados y procesados por un conjunto diverso de expertos en paralelo, y luego se reintegran sin problemas en la forma original del token. El mecanismo de múltiples cabezas permite que el modelo atienda colectivamente información de varios espacios de representación dentro de diferentes expertos, mientras que aumenta significativamente la activación de expertos, profundizando así la comprensión del contexto y mitigando el sobreajuste. Además, nuestro MH-MoE es sencillo de implementar y se desacopla de otros métodos de optimización de SMoE, lo que facilita su integración con otros modelos SMoE para mejorar el rendimiento. Los resultados experimentales extensos en tres tareas: modelado de lenguaje centrado en inglés, modelado de lenguaje multilingüe y tareas de modelado multimodal enmascarado, demuestran la efectividad de MH-MoE.

English

Sparse Mixtures of Experts (SMoE) scales model capacity without significant increases in training and inference costs, but exhibits the following two issues: (1) Low expert activation, where only a small subset of experts are activated for optimization. (2) Lacking fine-grained analytical capabilities for multiple semantic concepts within individual tokens. We propose Multi-Head Mixture-of-Experts (MH-MoE), which employs a multi-head mechanism to split each token into multiple sub-tokens. These sub-tokens are then assigned to and processed by a diverse set of experts in parallel, and seamlessly reintegrated into the original token form. The multi-head mechanism enables the model to collectively attend to information from various representation spaces within different experts, while significantly enhances expert activation, thus deepens context understanding and alleviate overfitting. Moreover, our MH-MoE is straightforward to implement and decouples from other SMoE optimization methods, making it easy to integrate with other SMoE models for enhanced performance. Extensive experimental results across three tasks: English-focused language modeling, Multi-lingual language modeling and Masked multi-modality modeling tasks, demonstrate the effectiveness of MH-MoE.