Branch-Train-MiX: Integración de Expertos LLM en un Modelo de Mezcla de Expertos LLM

Branch-Train-MiX: Mixing Expert LLMs into a Mixture-of-Experts LLM

March 12, 2024

Autores: Sainbayar Sukhbaatar, Olga Golovneva, Vasu Sharma, Hu Xu, Xi Victoria Lin, Baptiste Rozière, Jacob Kahn, Daniel Li, Wen-tau Yih, Jason Weston, Xian Li

cs.AI

Resumen

Investigamos métodos eficientes para entrenar Modelos de Lenguaje a Gran Escala (LLMs) con capacidades en múltiples dominios especializados, como programación, razonamiento matemático y conocimiento general. Nuestro método, denominado Branch-Train-MiX (BTX), comienza con un modelo base, el cual se bifurca para entrenar expertos de manera paralela y altamente eficiente, reduciendo los costos de comunicación. Una vez que los expertos individuales han sido entrenados de forma asíncrona, BTX combina sus parámetros de avance (feedforward) como expertos en capas de Mezcla de Expertos (Mixture-of-Expert, MoE) y promedia los parámetros restantes, seguido de una etapa de ajuste fino (finetuning) en MoE para aprender el enrutamiento a nivel de tokens. BTX generaliza dos casos especiales: el método Branch-Train-Merge, que carece de la etapa de ajuste fino en MoE para aprender el enrutamiento, y el reciclaje disperso (sparse upcycling), que omite la etapa de entrenamiento asíncrono de expertos. En comparación con enfoques alternativos, BTX logra el mejor equilibrio entre precisión y eficiencia.

English

We investigate efficient methods for training Large Language Models (LLMs) to possess capabilities in multiple specialized domains, such as coding, math reasoning and world knowledge. Our method, named Branch-Train-MiX (BTX), starts from a seed model, which is branched to train experts in embarrassingly parallel fashion with high throughput and reduced communication cost. After individual experts are asynchronously trained, BTX brings together their feedforward parameters as experts in Mixture-of-Expert (MoE) layers and averages the remaining parameters, followed by an MoE-finetuning stage to learn token-level routing. BTX generalizes two special cases, the Branch-Train-Merge method, which does not have the MoE finetuning stage to learn routing, and sparse upcycling, which omits the stage of training experts asynchronously. Compared to alternative approaches, BTX achieves the best accuracy-efficiency tradeoff.