AIO APEX

Mixture-of-Experts é a arquitetura que alimenta os maiores LLMs em produção — e funciona de forma diferente do que a maioria das pessoas imagina

Compartilhar:
Mixture-of-Experts é a arquitetura que alimenta os maiores LLMs em produção — e funciona de forma diferente do que a maioria das pessoas imagina

Quando a OpenAI lançou o GPT-4, a empresa se recusou a publicar a contagem de parâmetros. Meses depois, documentos vazados e benchmarks corroborantes sugeriram que ele usa uma arquitetura Mixture-of-Experts (MoE) com cerca de 1,8 trilhão de parâmetros totais distribuídos por oito sub-redes especialistas — mas ativa apenas cerca de 220 bilhões por forward pass. Essa única escolha de design explica tanto o teto de capacidade do modelo quanto sua economia de inferência de maneiras que uma contagem ingênua de parâmetros jamais conseguiria.

MoE é agora a arquitetura dominante para modelos de fronteira. O Google Gemini 1.5 usa MoE. Os modelos abertos Mixtral 8x7B e 8x22B da Mistral AI tornaram o MoE acessível para autohospedagem. A pesquisa interna da Meta sobre MoE para sucessores do Llama é bem documentada. Entender como ele realmente funciona — e onde ele genuinamente ajuda versus onde só melhora os slides — importa se você está decidindo quais modelos implantar ou como avaliar novos lançamentos.

A ideia central: computação condicional

Um transformer denso padrão como o Llama 2 70B ativa todos os seus 70 bilhões de parâmetros para cada token que processa. Isso é caro computacionalmente, mas previsível. O MoE substitui as feedforward layers (as camadas que compõem a maior parte da contagem de parâmetros de um transformer) por múltiplas redes "especialistas" paralelas mais um roteador leve. Para cada token, o roteador seleciona os top-k especialistas — tipicamente 2 entre 8 ou 16 — e apenas esses especialistas processam aquele token. Os resultados são ponderados e combinados.

A consequência prática: um modelo Mixtral 8x7B tem ~47 bilhões de parâmetros totais, mas cada token toca apenas cerca de 13 bilhões deles. Você obtém a maior parte da capacidade representacional de um modelo denso de 47B enquanto executa inferência a um custo mais próximo de 13B. A taxa de transferência (throughput) aproximadamente dobra em comparação com um modelo denso equivalente no mesmo hardware, para a mesma qualidade de saída.

O que o roteador realmente aprende

O roteador é uma pequena camada linear que produz uma distribuição de probabilidade sobre todos os especialistas disponíveis. Ele é treinado end-to-end com o restante do modelo usando gradiente descendente padrão — não há pré-treinamento separado ou rotulagem manual de qual especialista deve lidar com qual conteúdo. O que emerge é aproximadamente uma especialização por domínio: a análise dos padrões de roteamento do Mixtral mostra especialistas desenvolvendo preferências suaves para sintaxe de código, raciocínio em linguagem natural, recuperação factual, e assim por diante. Mas essa especialização é imprecisa e nem sempre está alinhada com intuições humanas sobre assunto.

Um problema persistente de engenharia é o balanceamento de carga (load balancing). Sem intervenção, o roteador tende a colapsar em um pequeno conjunto de especialistas "populares" e deixar os outros sem uso, desperdiçando capacidade. A correção padrão é uma perda auxiliar de balanceamento de carga adicionada durante o treinamento que penaliza a utilização desigual dos especialistas. Acertar a intensidade dessa perda é um hiperparâmetro que afeta significativamente tanto a qualidade do modelo quanto a eficiência do hardware — muito pouco, e ocorre colapso dos especialistas; muito, e o roteador não consegue aprender especialização significativa.

O gargalo de memória que o marketing ignora

Aqui é onde o MoE se torna complicado para quem implanta. Todos os parâmetros devem residir na memória mesmo que apenas uma fração ative por token. Um modelo Mixtral 8x22B — com ~141 bilhões de parâmetros totais — requer cerca de 280 GB de VRAM de GPU em precisão BF16 antes de considerar o KV cache. Isso significa pelo menos quatro GPUs H100 80GB apenas para armazenar os pesos, embora a taxa de transferência de inferência seja similar a um modelo denso muito menor.

Isso cria uma divisão de infraestrutura. Em um data center onde você pode dedicar um nó de 4 GPUs por réplica do modelo, o MoE é genuinamente mais barato por token. Em uma implantação onde você tenta colocalizar múltiplos modelos em hardware compartilhado, a pegada de memória do MoE o torna caro. É também por isso que a quantização importa mais para modelos MoE: levar o Mixtral 8x7B para precisão de 4 bits (aproximadamente 25 GB) é o que o torna prático para executar em uma única estação de trabalho de consumo ou em um servidor dual-GPU.

Paralelismo de especialistas como alavanca de escalabilidade

Para treinar modelos MoE muito grandes, uma técnica chamada expert parallelism distribui diferentes especialistas por diferentes GPUs físicas. Quando um token é roteado para o Especialista #5, a computação ocorre na GPU que contém os pesos do Especialista #5, e o resultado é enviado de volta. Isso transforma a comunicação all-reduce em transferências ponto a ponto mais localizadas e permite treinamento em escalas que, de outra forma, exigiriam muita memória por GPU.

O artigo do Switch Transformer do Google, de 2021, demonstrou isso com 1,6 trilhão de parâmetros — o primeiro modelo trilhão de parâmetros documentado publicamente. A descoberta principal: um MoE de 64 especialistas com o mesmo orçamento computacional que um modelo denso T5-XXL alcançou um ganho de 4x no tempo de treinamento, igualando ou superando a qualidade em benchmarks padrão. O artigo também documentou modos de falha: instabilidade de treinamento com altas contagens de especialistas, o problema de colapso de balanceamento de carga e sobrecarga de comunicação em configurações de múltiplos nós.

Onde o MoE genuinamente tem desempenho inferior aos modelos densos

O aprendizado few-shot em tarefas altamente específicas de domínio é uma área onde modelos MoE podem ter desempenho inferior a modelos densos de tamanho similar. Como o roteador atribui tokens probabilisticamente e diferentes tokens no mesmo prompt podem ir para diferentes especialistas, a "memória" do modelo do contexto inicial pode ser fragmentada entre especialistas de maneiras que prejudicam a coerência em documentos longos e especializados. Relatos anedóticos de implantações empresariais do Mixtral sugerem que modelos densos de custo de inferência equivalente às vezes produzem melhores resultados em textos jurídicos ou médicos onde a consistência exata da terminologia importa.

O tamanho do lote (batch size) também importa. A vantagem de throughput da arquitetura MoE é mais pronunciada em grandes tamanhos de lote, onde todos os especialistas obtêm utilização aproximadamente uniforme. No batch size 1 — um único usuário fazendo uma consulta em tempo real — você está ativando dois especialistas e esperando ociosos pelos outros seis. A latência por token pode ser pior do que a de um modelo denso com contagem equivalente de parâmetros ativados, devido à sobrecarga de roteamento. É por isso que as implantações em produção agrupam requisições agressivamente e por que os endpoints de API de streaming têm perfis de latência diferentes dos endpoints de inferência em lote.

Decisões práticas para equipes que avaliam modelos MoE

Se você está comparando um modelo denso de 70B com um modelo MoE como o Mixtral 8x22B para implantação, a comparação correta não é a contagem de parâmetros — é a pegada de memória versus a qualidade na sua carga de trabalho específica. Execute ambos na sua distribuição real de tarefas. O Mixtral 8x22B vencerá consistentemente o Llama 2 70B em benchmarks de raciocínio, mas a diferença diminui significativamente em tarefas estreitas de geração aumentada por recuperação (RAG) onde o conjunto de dados é homogêneo.

Para fine-tuning, modelos MoE apresentam um desafio particular: o fine-tuning LoRA aplicado apenas às camadas densas não tocará nos pesos dos especialistas, que contêm a maior parte do conhecimento especializado do modelo. O fine-tuning completo de modelos MoE consome muita memória. Variantes LoRA específicas para MoE que aplicam adaptadores às feedforward layers dos especialistas existem, mas ainda não são ferramentas padrão — verifique se seu framework de fine-tuning as suporta antes de se comprometer.

Os próprios pesos do roteador podem ser congelados durante o fine-tuning para preservar os padrões de especialização aprendidos durante o pré-treinamento. Isso funciona bem ao fazer fine-tuning para uma tarefa que está bem representada na distribuição original de treinamento. Ao adaptar para um domínio genuinamente novo, descongelar o roteador e aceitar um tempo maior de fine-tuning vale a pena.

O que vem a seguir

Direções de pesquisa ativamente exploradas incluem MoE esparso com mais de dois especialistas ativados por token (trocando computação por qualidade), roteamento hierárquico onde um roteador grosseiro seleciona "famílias" de especialistas antes de um roteador refinado selecionar especialistas específicos, e arquiteturas mixture-of-depths que roteiam tokens para diferentes camadas em vez de diferentes especialistas dentro de uma camada. O artigo de 2024 do Google DeepMind sobre mixture-of-depths mostrou que nem todo token precisa passar por todas as camadas do transformer, permitindo ganhos adicionais de computação condicional.

A lição arquitetural do MoE é consistente: as leis de escalabilidade recompensam a computação condicional. Gastar toda a sua computação em cada token para cada tarefa é um desperdício. Os modelos que importarão nos próximos dois anos serão cada vez mais sistemas híbridos que roteiam trabalho de forma inteligente — seja roteando para especialistas dentro de um modelo, roteando para diferentes modelos via orquestração, ou roteando para ferramentas externas. MoE é a primeira demonstração em escala de produção de que esse princípio funciona no nível dos pesos.

Compartilhar:
Mixture-of-Experts é a arquitetura que alimenta os maiores LLMs em produção — e funciona de forma diferente do que a maioria das pessoas imagina | AIO APEX