O Surgimento da Mistura de Especialistas para Modelos de Linguagem Grande e Eficientes

No mundo do processamento de linguagem natural (NLP), a busca por construir modelos de linguagem maiores e mais capazes tem sido uma força motriz por trás de muitos avanços recentes. No entanto, à medida que esses modelos crescem em tamanho, os requisitos computacionais para treinamento e inferência se tornam cada vez mais exigentes, pressionando os limites dos recursos de hardware disponíveis.

Então, entra a Mistura de Especialistas (MoE), uma técnica que promete aliviar essa carga computacional, permitindo o treinamento de modelos de linguagem maiores e mais poderosos. Abaixo, discutiremos MoE, exploraremos suas origens, funcionamento interno e suas aplicações em modelos de linguagem baseados em transformadores.

As Origens da Mistura de Especialistas

O conceito de Mistura de Especialistas (MoE) pode ser rastreado até os primeiros anos de 1990, quando os pesquisadores exploraram a ideia de computação condicional, onde partes de uma rede neural são ativadas seletivamente com base nos dados de entrada. Um dos trabalhos pioneiros nesse campo foi o paper “Adaptive Mixture of Local Experts” de Jacobs et al. em 1991, que propôs um quadro de aprendizado supervisionado para um ensemble de redes neurais, cada uma especializada em uma região diferente do espaço de entrada.

A ideia central por trás de MoE é ter várias “especialistas” redes, cada uma responsável por processar um subconjunto dos dados de entrada. Um mecanismo de controle, tipicamente uma rede neural em si, determina quais especialista(s) devem processar uma entrada dada. Essa abordagem permite que o modelo aloque seus recursos computacionais de forma mais eficiente, ativando apenas os especialistas relevantes para cada entrada, em vez de empregar a capacidade total do modelo para cada entrada.

Ao longo dos anos, vários pesquisadores exploraram e estenderam a ideia de computação condicional, levando a desenvolvimentos como MoEs hierárquicos, aproximações de baixa ordem para computação condicional e técnicas para estimar gradientes por meio de neurônios estocásticos e funções de ativação de limiar rígido.

Mistura de Especialistas em Transformadores

Mistura de Especialistas

Embora a ideia de MoE tenha estado por décadas, sua aplicação a modelos de linguagem baseados em transformadores é relativamente recente. Transformadores, que se tornaram o padrão de fato para modelos de linguagem de ponta, são compostos por várias camadas, cada uma contendo um mecanismo de autoatenção e uma rede neural feed-forward (FFN).

A inovação-chave na aplicação de MoE a transformadores é substituir as camadas FFN densas por camadas MoE esparsas, cada uma consistindo em várias especialistas FFN e um mecanismo de controle. O mecanismo de controle determina quais especialista(s) devem processar cada token de entrada, permitindo que o modelo ative seletivamente apenas um subconjunto de especialistas para uma sequência de entrada dada.

Um dos primeiros trabalhos que demonstrou o potencial de MoE em transformadores foi o paper “Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer” de Shazeer et al. em 2017. Esse trabalho introduziu o conceito de uma camada MoE com controle esparsamente acionado, que empregou um mecanismo de controle que adicionou esparsidade e ruído ao processo de seleção de especialistas, garantindo que apenas um subconjunto de especialistas fosse ativado para cada entrada.

Desde então, vários outros trabalhos avançaram ainda mais a aplicação de MoE a transformadores, abordando desafios como instabilidade de treinamento, balanceamento de carga e inferência eficiente. Exemplos notáveis incluem o Switch Transformer (Fedus et al., 2021), ST-MoE (Zoph et al., 2022) e GLaM (Du et al., 2022).

Benefícios da Mistura de Especialistas para Modelos de Linguagem

O benefício principal de empregar MoE em modelos de linguagem é a capacidade de aumentar o tamanho do modelo enquanto mantém um custo computacional relativamente constante durante a inferência. Ao ativar seletivamente apenas um subconjunto de especialistas para cada token de entrada, os modelos MoE podem alcançar o poder expressivo de modelos densos muito maiores enquanto requerem significativamente menos computação.

Por exemplo, considere um modelo de linguagem com uma camada FFN densa de 7 bilhões de parâmetros. Se substituirmos essa camada por uma camada MoE consistindo em oito especialistas, cada um com 7 bilhões de parâmetros, o número total de parâmetros aumenta para 56 bilhões. No entanto, durante a inferência, se ativarmos apenas dois especialistas por token, o custo computacional é equivalente ao de um modelo denso de 14 bilhões de parâmetros, pois computa duas multiplicações de matrizes de 7 bilhões de parâmetros.

Essa eficiência computacional durante a inferência é particularmente valiosa em cenários de implantação onde os recursos são limitados, como dispositivos móveis ou ambientes de computação de borda. Além disso, as reduzidas exigências computacionais durante o treinamento podem levar a substanciais economias de energia e uma pegada de carbono mais baixa, alinhando-se com a crescente ênfase em práticas de IA sustentáveis.

Desafios e Considerações

Embora os modelos MoE ofereçam benefícios convincentes, sua adoção e implantação também vêm com vários desafios e considerações:

1. Instabilidade de Treinamento: Os modelos MoE são conhecidos por serem mais propensos a instabilidades de treinamento em comparação com seus equivalentes densos. Esse problema surge da natureza esparsa e condicional das ativações de especialistas, o que pode levar a desafios na propagação de gradientes e convergência. Técnicas como a perda z do roteador (Zoph et al., 2022) foram propostas para mitigar essas instabilidades, mas mais pesquisas ainda são necessárias.
2. Ajuste Fino e Sobreajuste: Os modelos MoE tendem a sofrer de sobreajuste mais facilmente durante o ajuste fino, especialmente quando a tarefa downstream tem um conjunto de dados relativamente pequeno. Esse comportamento é atribuído à capacidade aumentada e à esparsidade dos modelos MoE, o que pode levar a uma especialização excessiva nos dados de treinamento. Estratégias de regularização e ajuste fino cuidadosas são necessárias para mitigar esse problema.
3. Requisitos de Memória: Embora os modelos MoE possam reduzir os custos computacionais durante a inferência, eles frequentemente têm requisitos de memória mais altos em comparação com modelos densos de tamanho semelhante. Isso ocorre porque todos os pesos dos especialistas precisam ser carregados na memória, mesmo que apenas um subconjunto seja ativado para cada entrada. Restrições de memória podem limitar a escalabilidade dos modelos MoE em dispositivos com recursos limitados.
4. Balanceamento de Carga: Para alcançar eficiência computacional ótima, é crucial balancear a carga entre os especialistas, garantindo que nenhum especialista esteja sobrecarregado enquanto outros permanecem subutilizados. Esse balanceamento de carga é tipicamente alcançado por meio de perdas auxiliares durante o treinamento e ajuste cuidadoso do fator de capacidade, que determina o número máximo de tokens que podem ser atribuídos a cada especialista.
5. Sobrecarga de Comunicação: Em cenários de treinamento e inferência distribuídos, os modelos MoE podem introduzir sobrecarga de comunicação adicional devido à necessidade de trocar informações de ativação e gradiente entre especialistas residentes em dispositivos ou aceleradores diferentes. Estratégias de comunicação eficientes e design de modelo consciente do hardware são essenciais para mitigar essa sobrecarga.

Apesar desses desafios, o potencial dos modelos MoE em permitir modelos de linguagem maiores e mais capazes tem impulsionado esforços significativos de pesquisa para abordar e mitigar essas questões.

Exemplo: Mixtral 8x7B e GLaM

Para ilustrar a aplicação prática de MoE em modelos de linguagem, consideremos dois exemplos notáveis: Mixtral 8x7B e GLaM.

Mixtral 8x7B é uma variante MoE do modelo de linguagem Mistral, desenvolvido pela Anthropic. Ele consiste em oito especialistas, cada um com 7 bilhões de parâmetros, resultando em um total de 56 bilhões de parâmetros. No entanto, durante a inferência, apenas dois especialistas são ativados por token, reduzindo efetivamente o custo computacional para o de um modelo denso de 14 bilhões de parâmetros.

Mixtral 8x7B demonstrou um desempenho impressionante, superando o modelo Llama de 70 bilhões de parâmetros enquanto oferece tempos de inferência muito mais rápidos. Uma versão ajustada para instruções de Mixtral 8x7B, chamada Mixtral-8x7B-Instruct-v0.1, também foi lançada, aprimorando ainda mais suas capacidades em seguir instruções de linguagem natural.

Outro exemplo notável é o GLaM (Google Language Model), um modelo MoE de grande escala desenvolvido pelo Google. O GLaM emprega uma arquitetura de transformador apenas decodificador e foi treinado em um conjunto de dados massivo de 1,6 trilhão de tokens. O modelo alcança um desempenho impressionante em avaliações de poucos e um único disparo, igualando a qualidade do GPT-3 enquanto usa apenas um terço da energia necessária para treinar o GPT-3.

O sucesso do GLaM pode ser atribuído à sua arquitetura MoE eficiente, que permitiu o treinamento de um modelo com um número vasto de parâmetros enquanto mantinha requisitos computacionais razoáveis. O modelo também demonstrou o potencial dos modelos MoE para serem mais eficientes em termos de energia e sustentáveis ambientalmente em comparação com seus equivalentes densos.

A Arquitetura Grok-1

GROK MISTURA DE ESPECIALISTAS

Grok-1 é um modelo MoE baseado em transformador com uma arquitetura única projetada para maximizar eficiência e desempenho. Vamos mergulhar nos principais especificações:

1. Parâmetros: Com um total impressionante de 314 bilhões de parâmetros, o Grok-1 é o maior modelo LLM aberto até o momento. No entanto, graças à arquitetura MoE, apenas 25% dos pesos (aproximadamente 86 bilhões de parâmetros) estão ativos a qualquer momento, melhorando as capacidades de processamento.
2. Arquitetura: O Grok-1 emprega uma arquitetura de Mistura de 8-Especialistas, com cada token sendo processado por dois especialistas durante a inferência.
3. Camadas: O modelo consiste em 64 camadas de transformador, cada uma incorporando atenção multi-cabeça e blocos densos.
4. Tokenização: O Grok-1 utiliza um tokenizador SentencePiece com um tamanho de vocabulário de 131.072 tokens.
5. Embebedos e Codificação Posicional: O modelo apresenta embebedos de 6.144 dimensões e emprega embebedos de posição rotatória, permitindo uma interpretação mais dinâmica dos dados em comparação com as codificações de posição fixas tradicionais.
6. Atenção: O Grok-1 usa 48 cabeças de atenção para consultas e 8 cabeças de atenção para chaves e valores, cada uma com um tamanho de 128.
7. Comprimento de Contexto: O modelo pode processar sequências de até 8.192 tokens de comprimento, utilizando precisão bfloat16 para computação eficiente.

Desempenho e Detalhes de Implementação

O Grok-1 demonstrou um desempenho impressionante, superando o LLaMa 2 70B e o Mixtral 8x7B com uma pontuação MMLU de 73%, mostrando sua eficiência e precisão em várias avaliações.

No entanto, é importante notar que o Grok-1 requer recursos de GPU significativos devido ao seu tamanho. A implementação atual na versão de código aberto se concentra em validar a correção do modelo e emprega uma implementação de camada MoE ineficiente para evitar a necessidade de kernels personalizados.

Não obstante, o modelo suporta particionamento de ativação e quantização de 8 bits, o que pode otimizar o desempenho e reduzir os requisitos de memória.

Em uma medida notável, xAI lançou o Grok-1 sob a licença Apache 2.0, tornando seus pesos e arquitetura acessíveis à comunidade global para uso e contribuições.

O lançamento de código aberto inclui um repositório de exemplo de código JAX que demonstra como carregar e executar o modelo Grok-1. Os usuários podem baixar os pesos do checkpoint usando um cliente de torrent ou diretamente por meio do HuggingFace Hub, facilitando o acesso a esse modelo inovador.

O Futuro da Mistura de Especialistas em Modelos de Linguagem

À medida que a demanda por modelos de linguagem maiores e mais capazes continua a crescer, a adoção de técnicas MoE é esperada para ganhar ainda mais impulso. Esforços de pesquisa em andamento estão focados em abordar os desafios restantes, como melhorar a estabilidade de treinamento, mitigar o sobreajuste durante o ajuste fino e otimizar os requisitos de memória e comunicação.

Uma direção promissora é a exploração de arquiteturas MoE hierárquicas, onde cada especialista é composto por vários subespecialistas. Essa abordagem pode potencialmente permitir uma escalabilidade e eficiência computacional ainda maiores, mantendo o poder expressivo de modelos grandes.

Além disso, o desenvolvimento de sistemas de hardware e software otimizados para modelos MoE é uma área ativa de pesquisa. Aceleradores especializados e frameworks de treinamento distribuído projetados para lidar eficientemente com os padrões de computação esparsos e condicionais dos modelos MoE podem melhorar ainda mais seu desempenho e escalabilidade.

Adicionalmente, a integração de técnicas MoE com outras inovações em modelagem de linguagem, como mecanismos de atenção esparsos, estratégias de tokenização eficientes e representações multimodais, pode levar a modelos de linguagem ainda mais poderosos e versáteis, capazes de lidar com uma ampla gama de tarefas.

Conclusão

A técnica de Mistura de Especialistas emergiu como uma ferramenta poderosa na busca por modelos de linguagem maiores e mais capazes. Ao ativar seletivamente especialistas com base nos dados de entrada, os modelos MoE oferecem uma solução promissora para os desafios computacionais associados ao aumento do tamanho dos modelos densos. Embora ainda haja desafios a serem superados, como instabilidade de treinamento, sobreajuste e requisitos de memória, os benefícios potenciais dos modelos MoE em termos de eficiência computacional, escalabilidade e sustentabilidade ambiental os tornam uma área emocionante de pesquisa e desenvolvimento.

À medida que o campo do processamento de linguagem natural continua a empurrar os limites do que é possível, a adoção de técnicas MoE provavelmente desempenhará um papel crucial na habilitação da próxima geração de modelos de linguagem. Ao combinar MoE com outras inovações em arquitetura de modelo, técnicas de treinamento e otimização de hardware, podemos esperar por modelos de linguagem ainda mais poderosos e versáteis que possam verdadeiramente entender e se comunicar com os humanos de maneira natural e transparente.