Branch-Train-MiX: Integrando Modelos de Linguagem Especializados em um Modelo de Linguagem de Mistura de Especialistas

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

Resumo

Investigamos métodos eficientes para treinar Modelos de Linguagem de Grande Escala (LLMs) a possuírem capacidades em múltiplos domínios especializados, como codificação, raciocínio matemático e conhecimento geral. Nosso método, denominado Branch-Train-MiX (BTX), começa com um modelo inicial, que é ramificado para treinar especialistas de forma embaraçosamente paralela, com alta taxa de transferência e custo de comunicação reduzido. Após os especialistas individuais serem treinados de forma assíncrona, o BTX combina seus parâmetros de propagação direta como especialistas em camadas de Mistura de Especialistas (MoE) e faz a média dos parâmetros restantes, seguido por uma etapa de ajuste fino MoE para aprender o roteamento em nível de token. O BTX generaliza dois casos especiais: o método Branch-Train-Merge, que não possui a etapa de ajuste fino MoE para aprender o roteamento, e o upcycling esparso, que omite a etapa de treinamento assíncrono de especialistas. Em comparação com abordagens alternativas, o BTX alcança o melhor equilíbrio entre precisão e eficiência.

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.