LocMoE+: Enrutador Mejorado con Consciencia de Características de Tokens para un Pre-Entrenamiento Eficiente de LLM

Las arquitecturas Mixture-of-Experts (MoE) han ganado recientemente popularidad en el ámbito de los grandes modelos de lenguaje (LLMs) debido a su capacidad para reducir significativamente los costos de entrenamiento e inferencia. Sin embargo, las arquitecturas MoE enfrentan desafíos, como disparidades significativas en el número de tokens asignados a cada experto y una tendencia a la homogeneización entre expertos, lo que afecta negativamente las capacidades de generación semántica del modelo. En este artículo, presentamos LocMoE+, una versión refinada del LocMoE de bajo costo, que incorpora las siguientes mejoras: (1) Cuantificación y definición de la afinidad entre expertos y tokens. (2) Implementación de una estrategia de enrutamiento adaptativo a nivel global para reorganizar tokens según sus puntuaciones de afinidad. (3) Reestimación del límite inferior para la capacidad de los expertos, que se ha demostrado que disminuye progresivamente conforme evoluciona la distribución de características de los tokens. Los resultados experimentales demuestran que, sin comprometer la convergencia ni eficacia del modelo, el número de tokens que cada experto procesa puede reducirse en más del 60%. Combinado con optimizaciones de comunicación, esto conduce a una mejora promedio en la eficiencia de entrenamiento que va desde el 5.4% hasta el 46.6%. Tras el ajuste fino, LocMoE+ muestra una mejora en el rendimiento de 9.7% a 14.1% en los conjuntos de datos GDAD, C-Eval y TeleQnA.