Lory: Mezcla de expertos completamente diferenciable para el preentrenamiento de modelos de lenguaje autorregresivos

Los modelos de mezcla de expertos (MoE) facilitan una escalabilidad eficiente; sin embargo, entrenar la red enrutadora conlleva el desafío de optimizar un objetivo discreto no diferenciable. Recientemente, se propuso una arquitectura MoE totalmente diferenciable, SMEAR (Muqeeth et al., 2023), que fusiona suavemente expertos en el espacio de parámetros; no obstante, su efectividad solo fue demostrada en el ajuste fino posterior en tareas de clasificación. En este artículo, presentamos Lory, el primer enfoque que escala tales arquitecturas al preentrenamiento de modelos de lenguaje autorregresivos. Lory introduce dos técnicas clave: (1) una estrategia de enrutamiento causal por segmentos que logra alta eficiencia en operaciones de fusión de expertos preservando la naturaleza autorregresiva de los modelos de lenguaje; (2) un método de agrupación de datos basado en similitud que promueve la especialización de expertos agrupando documentos similares en instancias de entrenamiento. Preentrenamos una serie de modelos Lory con 150 mil millones de tokens desde cero, con hasta 32 expertos y 30 mil millones (1.5 mil millones activos) de parámetros. Los resultados experimentales muestran mejoras significativas en desempeño respecto a modelos densos con parámetros equivalentes tanto en perplexidad (+13.9%) como en una variedad de tareas posteriores (+1.5%-11.1%). A pesar del enrutamiento a nivel de segmentos, los modelos Lory alcanzan un rendimiento competitivo comparado con modelos MoE de punta con enrutamiento a nivel de token. Además, demostramos que los expertos entrenados en Lory capturan especialización a nivel de dominio sin supervisión. Nuestro trabajo destaca el potencial de arquitecturas MoE totalmente diferenciables para el preentrenamiento de modelos de lenguaje y aboga por futuras investigaciones en esta área.