Branch-Train-MiX: Mezclando Expertos LLMs en un Mixture-of-Experts LLM

Investigamos métodos eficientes para entrenar Grandes Modelos de Lenguaje (LLMs) que posean capacidades en múltiples dominios especializados, como programación, razonamiento matemático y conocimiento mundial. Nuestro método, denominado Branch-Train-MiX (BTX), comienza a partir de un modelo semilla, que se bifurca para entrenar expertos de forma embarrassingly paralela con alta eficiencia y menor coste de comunicación. Después de entrenar individualmente a los expertos de manera asincrónica, BTX reúne sus parámetros feedforward como expertos en capas Mixture-of-Expert (MoE) y promedia los parámetros restantes, seguido de una etapa de ajuste fino MoE para aprender el enrutamiento a nivel de token. BTX generaliza dos casos especiales: el método Branch-Train-Merge, que no incluye la etapa de ajuste fino MoE para aprender el enrutamiento, y el reciclaje escaso (sparse upcycling), que omite la etapa de entrenamiento asincrónico de expertos. En comparación con enfoques alternativos, BTX logra el mejor equilibrio entre precisión y eficiencia.