Mixtral of Experts

Albert Q. Jiang; Alexandre Sablayrolles; Antoine Roux; Arthur Mensch; Blanche Savary; Chris Bamford; Devendra Singh Chaplot; Diego de las Casas; Emma Bou Hanna; Florian Bressand; Gianna Lengyel; Guillaume Bour; Guillaume Lample; Lélio Renard Lavaud; Lucile Saulnier; Marie-Anne Lachaux; Pierre Stock; Sandeep Subramanian; Sophia Yang; Szymon Antoniak; Teven Le Scao; Théophile Gervet; Thibaut Lavril; Thomas Wang; Timothée Lacroix; William El Sayed

Artigos · arXiv 2401.04088

Mixtral of Experts

Albert Q. Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, Gianna Lengyel, Guillaume Bour et al.

Publicado em 2024-01-08cs.LGcs.CL arXiv PDF

Ler com tradução e IA

Resumo em três linhas

🎯 Escalar a capacidade e o desempenho de transformers sem incorrer em latência de inferência e custos computacionais proibitivos.
🔧 Uma arquitetura decoder-only Sparse Mixture of Experts (SMoE) que substitui blocos feedforward padrão por 8 especialistas SwiGLU, selecionando 2 especialistas por token através de um roteador aprendido.
📈 Acessa 47B de parâmetros totais com apenas 13B ativos por token, superando o Llama 2 70B no MMLU (70,6% vs 69,9%) e GSM8K (74,4% vs 69,6%).

Resumo

Apresentamos o Mixtral 8x7B, um modelo de linguagem Sparse Mixture of Experts (SMoE). O Mixtral possui a mesma arquitetura do Mistral 7B, com a diferença de que cada camada é composta por 8 blocos feedforward (ou seja, especialistas). Para cada token, em cada camada, uma rede roteadora seleciona dois especialistas para processar o estado atual e combinar suas saídas. Embora cada token veja apenas dois especialistas, os especialistas selecionados podem ser diferentes a cada timestep. Como resultado, cada token tem acesso a 47B de parâmetros, mas utiliza apenas 13B de parâmetros ativos durante a inferência. O Mixtral foi treinado com um tamanho de contexto de 32k tokens e supera ou se iguala ao Llama 2 70B e ao GPT-3.5 em todos os benchmarks avaliados. Em particular, o Mixtral supera amplamente o Llama 2 70B em matemática, geração de código e benchmarks multilíngues. Também disponibilizamos um modelo ajustado para seguir instruções, o Mixtral 8x7B - Instruct, que supera o GPT-3.5 Turbo, Claude-2.1, Gemini Pro e o modelo Llama 2 70B - chat em benchmarks humanos. Tanto o modelo base quanto o instruct são lançados sob a licença Apache 2.0.

Resumo original (inglês)

We introduce Mixtral 8x7B, a Sparse Mixture of Experts (SMoE) language model. Mixtral has the same architecture as Mistral 7B, with the difference that each layer is composed of 8 feedforward blocks (i.e. experts). For every token, at each layer, a router network selects two experts to process the current state and combine their outputs. Even though each token only sees two experts, the selected experts can be different at each timestep. As a result, each token has access to 47B parameters, but only uses 13B active parameters during inference. Mixtral was trained with a context size of 32k tokens and it outperforms or matches Llama 2 70B and GPT-3.5 across all evaluated benchmarks. In particular, Mixtral vastly outperforms Llama 2 70B on mathematics, code generation, and multilingual benchmarks. We also provide a model fine-tuned to follow instructions, Mixtral 8x7B - Instruct, that surpasses GPT-3.5 Turbo, Claude-2.1, Gemini Pro, and Llama 2 70B - chat model on human benchmarks. Both the base and instruct models are released under the Apache 2.0 license.

Guia de leitura

Problema e motivação

Modelos de linguagem autorregressivos densos incorrem em custos computacionais diretamente proporcionais à sua contagem total de parâmetros para cada token processado. À medida que os modelos escalam para 70B parâmetros e além, servi-los com baixa latência e alto throughput torna-se computacionalmente custoso. Os autores têm como objetivo escalar a capacidade do modelo mantendo o custo de inferência por token comparável ao de modelos muito menores, criando um modelo de pesos abertos capaz de se igualar ou superar o Llama 2 70B e baselines proprietários como o GPT-3.5.

Ideias-chave

Sparse Mixture of Experts (SMoE): Substituição de blocos densos de redes feedforward (FFN) em todas as 32 camadas do transformer por 8 redes especialistas distintas, roteando cada token para apenas os 2 melhores especialistas (top-2).
Contagens de parâmetros desacopladas: A arquitetura contém 47B de parâmetros totais em suas camadas, mas executa apenas 13B de parâmetros ativos por passagem direta (forward pass) de token, correspondendo à pegada computacional de um modelo denso de tamanho moderado.
Janela de contexto completa de 32k: Suporte nativo para comprimentos de sequência de 32.768 tokens com 100% de acurácia de recuperação em benchmarks de passkey.
Roteamento sintático em vez de baseado em tópicos: A inspeção empírica do roteamento de tokens revela que a seleção de especialistas é orientada principalmente pela função sintática e localidade temporal, em vez de domínios temáticos de alto nível.

Método (com as equações principais)

O Mixtral modifica a arquitetura do Mistral 7B ( $32$ camadas, dimensão oculta $4096$ , $32$ cabeças de atenção, $8$ cabeças de chave-valor) substituindo cada sub-bloco FFN por um módulo MoE composto por $n = 8$ especialistas ( $E_{0}, \dots, E_{n - 1}$ ), cada um implementado como uma rede SwiGLU com dimensão intermediária de $14336$ .

Para uma representação de token de entrada $x$ , uma rede roteadora calcula uma distribuição esparsa de gating entre os especialistas usando os logits top- $K$ ( $K = 2$ ): $G (x) := Softmax (TopK (x \cdot W_{g}))$ onde o operador Top-K atribui $- \infty$ a todas as coordenadas fora dos maiores $K$ valores de logits $ℓ = x \cdot W_{g} \in R^{n}$ : $(TopK (ℓ))_{i} := {ℓ_{i} - \infty se ℓ_{i} estiver entre as top- K coordenadas de ℓ, caso contr \overset{a}{ˊ} rio.$

A saída geral da camada $y$ é calculada como a soma ponderada pelo gating sobre as saídas dos especialistas: $y = i = 0 \sum n - 1 G (x)_{i} \cdot SwiGLU_{i} (x)$ A execução eficiente é alcançada por meio de kernels como o MegaBlocks, que representam a computação MoE como multiplicações de matrizes esparsas, juntamente com o Expert Parallelism através de múltiplas GPUs.

Resultados (com os números mais importantes)

Benchmarks gerais e de raciocínio: No MMLU (5-shot), o Mixtral 8x7B atinge 70,6% em comparação a 69,9% do Llama 2 70B e 70,0% do GPT-3.5. No GSM8K (8-shot, maj@8), atinge 74,4% contra 69,6% do Llama 2 70B. No benchmark MATH (4-shot, maj@4), o Mixtral alcança 28,4% contra 13,8% do Llama 2 70B.
Geração de código: O Mixtral supera o Llama 2 70B de forma substancial no HumanEval (40,2% vs 29,3%) e no MBPP (60,7% vs 49,8%).
Tarefas multilíngues: A sobreamostragem (upsampling) de dados multilíngues permite que o Mixtral supere o Llama 2 70B em francês, alemão, espanhol e italiano no ARC-Challenge, HellaSwag e MMLU (por exemplo, MMLU em francês: 70,9% vs 64,3%).
Fine-Tuning de instruções e Arena Elo: O Mixtral 8x7B – Instruct (treinado com SFT e Direct Preference Optimization) obtém pontuação 8,30 no MT-Bench e alcançou um Elo de 1121 no LMSys Chatbot Arena, superando o Claude-2.1 (1117), GPT-3.5-Turbo (1117) e Llama-2-70b-chat (1077).
Recuperação de contexto: Atinge 100% de acurácia de recuperação na tarefa sintética de passkey retrieval em toda a janela de contexto de 32k.

Limitações e questões em aberto

Pegada de memória vs. computação: Embora o custo computacional escale com os 13B parâmetros ativos, o deployment exige VRAM suficiente para armazenar todos os 47B parâmetros, exigindo configurações multi-GPU ou técnicas de offloading.
Gargalos no Expert Parallelism: O artigo descobre uma alta localidade temporal — tokens consecutivos compartilham as mesmas atribuições de especialistas com muito mais frequência do que por acaso (até 28,4% de repetição de primeira escolha na camada 15 vs 12,5% aleatório). Isso pode causar desbalanceamento de carga e sobrecarga de GPUs específicas sob Expert Parallelism.
Falta de especialização por tópicos: A seleção do roteador não particiona naturalmente por domínios conceituais (por exemplo, biologia vs. filosofia), deixando questões em aberto sobre como estruturar o gating caso a especialização de especialistas em domínios específicos seja explicitamente desejada.

Termos-chave

Sparse Mixture of Experts (SMoE): Uma arquitetura de rede neural onde a computação condicional encaminha cada token de entrada para um subconjunto esparso de sub-redes feedforward (especialistas) em vez de ativar todos os parâmetros.
Active Parameter Count: O número de parâmetros efetivamente carregados e computados durante a inferência para um único token, que é de 13B no Mixtral em comparação com seus 47B parâmetros esparsos totais.
SwiGLU: Uma variante de ativação de unidade linear com gating que combina Swish e multiplicação elemento a elemento, usada como a arquitetura interna de cada bloco feedforward especialista.
TopK Gating: Um mecanismo de roteamento que aplica uma projeção linear à representação do token, retém os K maiores valores enquanto mascara o restante para menos infinito e os normaliza via softmax.
Expert Parallelism (EP): Uma estratégia de particionamento distribuído que mapeia diferentes especialistas para GPUs distintas e roteia as ativações de tokens entre os dispositivos para o seu especialista designado.
Direct Preference Optimization (DPO): Um algoritmo de alinhamento que otimiza os parâmetros do modelo de linguagem diretamente em pares de preferência, sem treinar um modelo explícito de recompensa por aprendizado por reforço.
Passkey Retrieval: Um teste de avaliação sintético que mede a capacidade de um modelo de localizar e retornar uma chave curta posicionada aleatoriamente dentro de prompts de contexto longo de até 32k tokens.
Positional Locality: A tendência observada onde tokens consecutivos em uma sequência são repetidamente atribuídos ao mesmo especialista muito acima da probabilidade aleatória uniforme.

Perguntas

Por que o Mixtral é chamado de 8x7B se ele tem apenas 47B parâmetros em vez de 56B?

Embora existam 8 especialistas por camada e a arquitetura seja derivada do Mistral 7B, os pesos de autoatenção e as camadas de embedding são compartilhados entre todos os especialistas. Apenas os sub-blocos feedforward são replicados 8 vezes, resultando em 47B de parâmetros totais em vez de 8 × 7B = 56B.

Quantos especialistas são ativados para cada token durante a inferência?

O roteador seleciona exatamente 2 especialistas ( $K = 2$ ) dos 8 especialistas disponíveis em cada camada para cada token, resultando em 13B parâmetros ativos por passagem direta (forward pass).

Especialistas diferentes se especializam em tópicos acadêmicos ou domínios distintos?

Nenhuma especialização óbvia por domínio foi encontrada. A análise no The Pile mostrou distribuições de atribuição de especialistas praticamente idênticas em temas como ArXiv, PubMed e PhilPapers; em vez disso, as decisões de roteamento correlacionam-se fortemente com a sintaxe, a identidade do token e a localidade posicional.

Sob qual licença de software os modelos Mixtral 8x7B são disponibilizados?

Tanto o modelo base Mixtral 8x7B quanto o Mixtral 8x7B – Instruct são lançados sob a permissiva licença Apache 2.0 para uso acadêmico e comercial.

A explicação desta página foi escrita por IA a partir do artigo e pode conter erros. Consulte o artigo para os detalhes.

Ler com tradução e IA