Branch-Train-MiX: Misturando LLMs Especializados em um LLM Mixture-of-Experts

Investigamos métodos eficientes para treinar Grandes Modelos de Linguagem (LLMs) para possuírem habilidades em múltiplos domínios especializados, como codificação, raciocínio matemático e conhecimento geral. Nosso método, chamado Branch-Train-MiX (BTX), parte de um modelo semente, que é ramificado para treinar especialistas de forma paralela e embarrassingly paralela, com alta taxa de processamento e custo de comunicação reduzido. Após os especialistas serem treinados de forma assíncrona, o BTX reúne seus parâmetros feedforward como especialistas em camadas Mixture-of-Expert (MoE) e faz a média dos parâmetros restantes, seguida por uma etapa de fine-tuning MoE para aprender o roteamento ao nível de tokens. O BTX generaliza dois casos especiais: o método Branch-Train-Merge, que não possui a etapa de fine-tuning MoE para aprender roteamento, e o upcycling esparso, que omite a etapa de treinamento assíncrono dos especialistas. Em comparação com abordagens alternativas, o BTX alcança a melhor relação acurácia-eficiência.