O que é combinação de especialistas?

As arquiteturas de combinação de especialistas permitem que modelos em grande escala, mesmo aqueles com muitos bilhões de parâmetros, reduzam consideravelmente os custos de computação durante o treinamento prévio e alcancem um desempenho mais rápido durante o tempo de inferência. Em termos gerais, ela alcança essa eficiência ativando de forma seletiva somente os especialistas específicos necessários para uma determinada tarefa, em vez de ativar toda a rede neural em todas as tarefas.

Embora grande parte da implementação moderna de configurações de combinação de especialistas tenha sido desenvolvida ao longo (aproximadamente) da última década, a premissa central por trás dos modelos de MoE tem origem no artigo de 1991 "Adaptive Mixture of Local Experts". O artigo propôs treinar um sistema de IA composto por redes separadas, cada uma especializada em um subconjunto diferente de casos de treinamento. Isso envolveu o treinamento das próprias "redes de especialistas" e de uma rede de gating que determina qual especialista deve ser usado para cada subtarefa. Quando comparado a um modelo convencional análogo, os autores descobriram que seu modelo experimental era significativamente mais rápido de treinar: atingiu o limite de precisão alvo na metade do tempo de treinamento do modelo convencional.¹

Nos últimos anos, à medida que os principais modelos de deep learning utilizados em IA generativa se tornaram cada vez maiores e mais exigentes em termos de computação, a combinação de especialistas oferece uma forma de lidar com a troca entre a maior capacidade de modelos maiores e a maior eficiência de modelos menores. Isso foi explorado de forma mais notável no campo do processamento de linguagem natural (PLN): alguns dos principais grandes modelos de linguagem (LLMs), como o Mixtral 8x7B da Mistral e (de acordo com alguns relatórios) o GPT-4 da OpenAI,² empregaram a arquitetura MoE.

Os modelos modernos de deep learning são criados a partir de redes neurais artificiais formadas por várias camadas de nós interconectados (ou "neurônios"). Cada neurônio tem uma função de ativação: uma operação matemática realizada nos dados recebidos da camada anterior, cuja produção informa o input alimentado à camada seguinte. As redes neurais feed-forward (FFNs) clássicas processam as informações passando progressivamente os dados de input dos neurônios em uma camada para os neurônios na camada seguinte até chegarem a uma camada externa onde ocorrem as previsões finais. Algumas arquiteturas de redes neurais incorporam elementos adicionais, como os mecanismos de autoatenção dos modelos transformadores que detectam padrões e dependências adicionais nos dados de input.

As conexões entre diversas camadas e neurônios são mediadas por parâmetros do modelo que podem ser aprendidos: pesos (e vieses) variáveis que ampliam ou reduzem a influência que uma determinada parte da produção da rede tem sobre outras partes. Um modelo de deep learning "aprende" ajustando esses parâmetros, utilizando algoritmos de otimização como gradiente descendente de uma forma que aumenta a precisão das previsões.

Embora um número maior de parâmetros aumente a capacidade do modelo, sua capacidade de absorver informações e padrões nele contidos, ele também eleva os recursos computacionais necessários para treinar e operar o modelo. Em um modelo típico de deep learning, o que nesse contexto é conhecido como modelo denso, toda a rede é executada para processar todo e qualquer input. Isso cria um equilíbrio entre a capacidade do modelo e a praticidade.

Ao contrário dos modelos densos convencionais, a combinação de especialistas usa computação condicional para impor a escassez: em vez de usar toda a rede para cada input, os modelos MoE aprendem uma função de mapeamento computacionalmente barata que determina quais partes da rede – em outras palavras, quais especialistas – são mais eficaz para processar um determinado input, como um único token usado para representar uma palavra ou fragmento de palavra em tarefas de PLN.

Isso permite que a capacidade do modelo seja aumentada (expandindo o número total de parâmetros) sem um aumento correspondente na carga computacional necessária para treiná-lo e executá-lo (porque nem todos esses parâmetros serão necessariamente usados em um determinado momento).

Os modelos MoE processam dados designando vários "especialistas", cada um com sua própria subrede dentro de uma rede neural maior, e treinando uma rede de gating (ou roteador) para ativar somente os especialistas específicos mais adequados a um determinado input.

O principal benefício da abordagem de MoE é que, impondo a escassez em vez de ativar toda a rede neural para cada token de input, a capacidade do modelo pode ser aumentada, mantendo basicamente os custos computacionais constantes.

Em um nível de arquitetura, isso é alcançado substituindo camadas tradicionais de rede feed-forward (FFN) densas por camadas (ou blocos) esparsas de MoE. No jargão das redes neurais, "bloco" refere-se a um elemento estrutural recorrente que executa uma função específica. Em um modelo de MoE esparso (SMoE), esses blocos de especialistas podem ser camadas únicas, FFNs autônomos ou inclusive MoEs aninhados em si mesmos.

Por exemplo, no modelo de linguagem Mixtral 8x7B da Mistral, cada camada é composta por 8 blocos feed-forward, ou seja, especialistas, cada um com 7 bilhões de parâmetros. Para cada token, em cada camada, uma rede de roteadores seleciona dois desses oito especialistas para processar os dados. Em seguida, combina os resultados desses dois especialistas e passa o resultado para a camada seguinte. Os especialistas específicos selecionados pelo roteador em uma determinada camada podem ser especialistas diferentes daqueles selecionados na camada anterior ou seguinte.³

Um SMoE pode ser composto inteiramente por camadas esparsas de MoE, mas muitas arquiteturas de modelo MoE envolvem blocos esparsos e densos. Isso também se aplica ao Mixtral, no qual os blocos responsáveis pelo mecanismo de autoatenção do modelo são compartilhados entre todos os 8 especialistas. Na prática, isso torna designações como "8x7B" potencialmente enganosas: como muitos dos parâmetros do modelo são compartilhados por cada subrede de especialistas de 7 bilhões de parâmetros, o Mixtral tem um total de cerca de 47 bilhões de parâmetros, não 56 bilhões, como se poderia imaginar por meio de uma multiplicação simples.

Essa contagem geral de parâmetros é comumente chamada de contagem de parâmetros esparsos e geralmente pode ser entendida como uma medida da capacidade do modelo. O número de parâmetros que serão realmente usados para processar um token (à medida que transita por alguns blocos de especialistas e ignora outros) é chamado de contagem de parâmetros ativos e pode ser entendido como uma medida dos custos computacionais do modelo. Embora cada input de token para o Mixtral tenha acesso a 46,7 bilhões de parâmetros, apenas 12,9 bilhões de parâmetros ativos são usados para processar um determinado exemplo.

É fundamental conhecer essa utilização ideal de contagens de parâmetros para entender o lado positivo dos modelos MoE. Por exemplo, o Mixtral supera a variante de 70 bilhões de parâmetros do Llama 2 do Meta na maioria dos benchmarks, com velocidade muito maior, embora tenha um terço a menos de parâmetros totais e utilize menos de 20% dos parâmetros ativos no momento da inferência.³

No entanto, vale a pena observar que a contagem geral de parâmetros de uma MoE esparsa não é totalmente irrelevante para os requisitos computacionais. Apesar de usar apenas um subconjunto de parâmetros durante a inferência, todos os parâmetros do modelo devem ser carregados na memória, o que significa que a eficiência computacional desfrutada pelos SMOEs em muitos aspectos não se aplica aos seus requisitos de RAM/VRAM.

A chave para o conceito (e a eficiência) das MoEs é que somente alguns dos especialistas (e, portanto, parâmetros) em uma camada esparsa serão ativados em um determinado momento, reduzindo assim os requisitos computacionais ativos.

Embora a computação condicional tenha sido proposta há muito tempo como um meio teórico de dissociar as demandas computacionais do aumento da capacidade do modelo, os desafios algorítmicos e de desempenho para sua execução bem-sucedida não foram superados até o artigo de 2017 de Shazeer et al, "Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer."⁴

As vantagens de camadas esparsas em relação a camadas densas são mais evidentes ao lidar com dados de alta dimensão, onde os padrões e dependências são muitas vezes complexos e não lineares: por exemplo, em tarefas de PLN que exigem que um modelo processe uma longa sequência de texto, cada palavra normalmente está relacionada apenas a um pequeno subconjunto de outras palavras nessa sequência. Isso torna os SMoEs uma área de enorme potencial no campo de LLMs, onde modelos de MoE bem calibrados podem aproveitar os benefícios da escassez sem sacrificar o desempenho. Modelos de MoE com portas esparsas também foram aplicados com sucesso a tarefas de visão computacional^{5 6} e continuam sendo uma área de estudo ativa nesse campo. 

Esta dispersão é alcançada por meio da computação condicional: a ativação dinâmica de parâmetros específicos em resposta a inputs específicos. O desenho eficaz da rede de gating (ou "roteador"), que impõe essa computação condicional é, portanto, essencial para o sucesso dos modelos MoE.

Vários mecanismos de gating podem ser usados para selecionar quais especialistas são utilizados em uma determinada situação. A função de gating correta é crítica para o desempenho do modelo, pois uma estratégia de roteamento ruim pode resultar em alguns especialistas sendo mal treinados ou excessivamente especializados e reduzir a eficácia de toda a rede.

Um mecanismo de gating típico em uma configuração tradicional de MoE, apresentado no artigo seminal de Shazeer, usa a função softmax: para cada um dos especialistas, por exemplo, o roteador prevê um valor de probabilidade (com base nos pesos das conexões desse especialista para o parâmetro atual) desse especialista produzindo o melhor resultado para um determinado input; em vez de calcular o resultado de todos os especialistas, o roteador calcula apenas o resultado (o que ele prevê ser) dos top k especialistas para aquele exemplo. Conforme descrito anteriormente, a Mixtral usa essa estratégia clássica de roteamento top-k: especificamente, eles usam o roteamento top-2, ou seja, k=2, selecionando os 2 melhores (de um total de 8) especialistas.

Em seu influente artigo de 2021, "Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficiency Sparsity", Fedus et al. levaram o roteamento top-k ao extremo: trabalhando com o T5 LLM do Google, eles substituíram as camadas FFN do modelo por 128 especialistas e implementaram k=1, também chamado de "roteamento rígido". Mesmo ao escalar o modelo para até um trilhão de parâmetros, essa configuração melhorou as velocidades de pré-treinamento em 400%.⁶

Apesar de seus muitos benefícios, as MoEs adicionam uma complexidade significativa ao processo de treinamento. Uma desvantagem importante da estratégia de roteamento top-k "vanilla" é o potencial da rede de gating convergir para a ativação de apenas alguns especialistas. Este é um problema que se reforça por si só: se alguns poucos especialistas forem selecionados de forma desproporcional desde o início, esses especialistas serão treinados mais rapidamente e, então, continuarão sendo selecionados mais vezes, pois agora produzem previsões mais confiáveis do que os outros especialistas menos treinados. Essa carga desequilibrada significa que outros especialistas acabam, figurativa e literalmente, como peso(s) morto(s).

Para mitigar isso, Shazeer et al apresentaram o top-k gating ruidoso: é adicionado um ruído gaussiano aos valores de probabilidade previstos para cada especialista, introduzindo um grau de aleatoriedade que incentiva uma ativação dos especialistas distribuída com mais uniformidade. Acrescentaram também dois termos de regularização que podem ser treinados à seleção de especialistas: a minimização da perda de balanceamento de carga penaliza a dependência excessiva de um único especialista, enquanto a minimização da perda de diversidade de especialistas recompensa a utilização equilibrada de todos os especialistas.

O artigo de 2020 do Google, "GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding", introduziu dois meios adicionais de balanceamento de carga:

• Roteamento aleatório: enquanto o especialista "top" em sua configuração top-2 é selecionado usando a função softmax padrão, o segundo especialista é escolhido de forma semi-aleatória (com a probabilidade de qualquer especialista ser escolhido proporcional ao peso de sua conexão). O segundo especialista mais bem classificado é, portanto, mais provável de ser selecionado, mas não é mais garantido que seja selecionado.

• Capacidade do especialista: os autores definem um limite que define o número máximo de tokens que podem ser processados por qualquer especialista. Se algum dos top-2 especialistas escolhidos estiver com a capacidade esgotada, o token será considerado "transbordado" e passa para a camada seguinte da rede.⁷

Como mencionado anteriormente, as vantagens dos modelos esparsos às vezes são moderadas por sua complexidade adicional. As dificuldades de implementação de MoEs ficam particularmente evidentes no processo de ajuste fino . Os modelos esparsos são mais propensos ao sobreajuste do que os modelos densos tradicionais e a presença de camadas MoE esparsas e camadas FFN densas complica o uso de uma abordagem única para todos.

Uma série de observações e abordagens têm sido propostas para reduzir a instabilidade ao ajustar as MoEs. Os autores do artigo "Switch Transformers" observaram que as variantes com menos especialistas tiveram um ajuste fino mais bem-sucedido, o que sugere que os benefícios de um número maior de especialistas em pré-treinamento podem ser contrabalançados por seu obstáculo à especialização em tarefas posteriores.

No artigo de 2022 "ST-MoE: Designing Stable and Transferable Sparse Expert Models", Zoph et al compararam os resultados de 5 abordagens diferentes: ajuste fino de todos os parâmetros ("Todos"), apenas parâmetros não MoE ("Não MoE" ), apenas parâmetros MoE ("MoE"), apenas os parâmetros de autoatenção do codificador-decodificador ("Atenção") e apenas os parâmetros FFN não MoE ("FFN").

• Quase nenhuma diferença foi encontrada entre Todos e Não MoE

• O ajuste fino apenas dos parâmetros de atenção resultou em uma pequena redução no desempenho.

• A atualização apenas dos parâmetros da MoE degradou significativamente o desempenho do modelo, apesar do fato de que cerca de 80% dos parâmetros do modelo residiam nas camadas esparsas da MoE.

• O FFN foi a única abordagem que melhorou o desempenho em relação à linha de base Todos.

Os autores levantaram a hipótese de que, como as camadas de especialistas representavam apenas um quarto do total de camadas de seu modelo, e um token verá no máximo apenas dois especialistas por camada, isolar os parâmetros de MoE resulta em atualizações de peso menos abrangentes (e, portanto, maior sobreajuste e perda de treinamento).⁸

Um artigo de julho de 2023, "Mixture-of-Experts Meets Instruction Tuning", explorou o impacto do ajuste de instruções nos modelos MoE com equivalentes do T5 e Flan-T5 do Google, uma versão do T5 ajustada com o protocolo Flan do Google - LLMs como linha de base. Seu experimento comparou quatro configurações: ajuste fino de um modelo T5 denso, ajuste fino de um modelo Flan-T5 denso, ajuste fino de um modelo MoE e ajuste fino de um modelo Flan-MoE ajustado por instrução.

Como esperado, o equivalente denso do T5 superou o MoE após o ajuste fino. Mas por outro lado, o modelo Flan-MoE com ajuste fino superou consideravelmente o modelo Flan-T5 ajustado. Além disso, a melhora do Flan-MoE em relação ao MoE foi ainda maior do que a melhora do Flan-T5 em relação ao T5.9 original

De forma promissora, isso sugere que, apesar de suas dificuldades com o ajuste fino padrão, os modelos MoE na verdade se beneficiam mais do ajuste de instruções do que seus equivalentes densos. Essa descoberta foi feita com o aclamado lançamento do Mixtral 8x7B Instruct, variante ajustada por instruções do Mixtral oferecida como modelo de base no IBM watsonx.ai.