Lory: Mistura de Experts totalmente diferenciável para pré-treinamento de modelo de linguagem autoregressivo

Modelos de mistura de experts (MoE) facilitam a escalabilidade eficiente; no entanto, treinar a rede roteadora introduz o desafio de otimizar um objetivo discreto e não diferenciável. Recentemente, uma arquitetura MoE totalmente diferenciável, SMEAR, foi proposta (Muqeeth et al., 2023), que mescla suavemente os experts no espaço dos parâmetros; entretanto, sua eficácia foi demonstrada apenas em fine-tuning downstream em tarefas de classificação. Neste artigo, apresentamos o Lory, a primeira abordagem que escala tais arquiteturas para pré-treinamento de modelo de linguagem autoregressivo. Lory introduz duas técnicas chave: (1) uma estratégia de roteamento causal por segmentos que alcança alta eficiência para operações de mesclagem de experts preservando a natureza autoregressiva dos modelos de linguagem; (2) um método de agrupamento de dados baseado em similaridade que incentiva a especialização dos experts agrupando documentos similares nas instâncias de treino. Pré-treinamos uma série de modelos Lory em 150B tokens do zero, com até 32 experts e 30B (1.5B ativos) parâmetros. Resultados experimentais mostram ganhos significativos de desempenho sobre modelos densos equivalentes em números de parâmetros tanto em perplexidade (+13.9%) quanto em diversas tarefas downstream (+1.5%-11.1%). Apesar do roteamento a nível de segmento, modelos Lory alcançam desempenho competitivo comparado a modelos MoE de última geração com roteamento a nível de token. Demonstramos ainda que os experts treinados no Lory capturam especialização a nível de domínio sem supervisão. Nosso trabalho destaca o potencial das arquiteturas MoE totalmente diferenciáveis para pré-treinamento de modelos de linguagem e incentiva pesquisas futuras nessa área.