O roteamento em Mistura de Especialistas (MoE) é um mecanismo central em arquiteturas modernas de Machine learning que atribui dinamicamente cada token de entrada a um subconjunto de redes especializadas. Em vez de ativar todos os parâmetros para cada entrada, um roteador, ou rede de portão, aprende a direcionar tokens aos especialistas mais relevantes, equilibrando capacidade do modelo com eficiência computacional. Essa abordagem permite que modelos de linguagem de grande escala cresçam para trilhões de parâmetros mantendo os custos de inferência gerenciáveis, pois apenas uma fração da rede é ativada por token.
O conceito origina-se dos trabalhos do início dos anos 1990 sobre misturas adaptativas de especialistas locais, mas ganhou destaque nos anos 2010 com a introdução de camadas MoE com portas esparsas em aprendizado profundo. Implementações modernas, como as usadas em modelos da Google DeepMind e OpenAI, dependem de funções de roteamento aprendidas que geram distribuições de probabilidade sobre os especialistas, muitas vezes incorporando perdas auxiliares para garantir uso equilibrado. O roteamento tornou-se uma escolha de design crítica em sistemas de IA generativa de última geração, influenciando tanto a eficiência do treinamento quanto a qualidade final do modelo.
Desenvolvimento Histórico
A ideia de rotear entradas para componentes especializados predata o aprendizado profundo, com raízes em métodos de conjunto e redes neurais modulares. Em 1991, Michael Jordan e Robert Jacobs publicaram um trabalho fundamental sobre misturas hierárquicas de especialistas, onde uma rede de portão ponderava as saídas de várias redes feedforward. Esse arcabouço inicial estabeleceu o princípio central: aprender a particionar o espaço de entrada para que diferentes especialistas lidassem com diferentes regiões.
Em 2017, pesquisadores da Google DeepMind (então Google Brain) introduziram a camada MoE de portas esparsas em um artigo liderado por Noam Shazeer. Esse design usava uma função de portão softmax para selecionar os principais k especialistas para cada token, com k geralmente definido em 1 ou 2. Os autores demonstraram que camadas MoE poderiam escalar para centenas de especialistas, alcançando resultados de ponta em benchmarks de modelagem de linguagem e tradução automática, reduzindo o custo computacional por exemplo. Esse trabalho influenciou diretamente arquiteturas posteriores como o Switch Transformer (2021), que simplificou o roteamento para selecionar apenas um especialista por token, e o arcabouço GShard para treinamento massivamente paralelo.
Mecanismos de Roteamento
Funções de roteamento podem ser categorizadas em vários tipos. O mais comum é o roteamento por escolha de token, onde cada token seleciona independentemente seus principais k especialistas baseado em um escore aprendido. Isso é implementado como uma projeção linear do estado oculto do token seguida de um softmax sobre os índices dos especialistas. Os especialistas selecionados então em solo processam o token, e suas saídas são ponderadas pelas probabilidades de roteamento e somadas.
Uma alternativa é o roteamento por escolha de especialista, onde cada especialista seleciona os top k tokens de um lote, garantindo que todos receberam uma carga mínima. Essa abordagem, popularizada no Switch Transformer e posteriormente refinada em modelos como o Mixtral, resolve tecnicamente problemas de desbalanceamento de carga, mas requer um cuidadoso manuseio das atribuições de token a especialista.
Outras variantes incluem roteamento baseado em hash, que usa funções de hash determinísticas para atribuir tokens sem falhas em создания learned, e roteamento hierárquico, onde um sistema de portão de dois níveis primeiro seleciona um grupo de especialistas e, então, um especialista específico dentro desse grupo. Cada método executa trocas entre flexibilidade, sobrecarga computacional e estabilidade de treinamento.
Balanceamento de Carga e Perdas Auxiliares
Um grande desafio no roteamento MoE é o desbalanceamento de carga: o roteador pode aprender a sempre selecionar alguns especialistas populares, deixando outros subtreinados. Para mitigar isso, a maioria das implementações adiciona uma perda auxiliar de balanceamento de carga que punementa a distribuição desigual de tokens entre especialistas. A formulação padrão, introduzida no artigo de Shazeer de 2017, calcula a fração de tokens roteados para cada especialista e constante uniformidade por meio de um termo de entropia cruzada.
Abordagens mais recentes, como a perda de balanceamento de carga do Switch Transformer, usam um coeficiente mais simples que multiplica o produto escalar entre a célula média do roteador e a clicle實 count de tokens por especialista. Alguns sistemas, como GShard, empregam fatores de capacidade que limitam quantos tokens cada especialista pode processar, forçando o roteador a distribuir o trabalho. Essas técnicas são essenciais para treinamento para estáveis e prevenindo o colapso do especialista, onde alguns se tornam pesos mortos.
Integração Arquitetural
O roteamento é tipicamente aplicado dentro de blocos de transformações, substituindo a rede de feedforward (FFN) por uma camada MoE. Em um transformador padrão, cada token passa por uma subcamada de atenção multi-cabeça seguida por uma FFN posição-savvy. Em um transformador MoE, a FFN é substituída por um conjunto de FFNs de especialistas, cada uma com seus próprios parâmetros, e um roteador que seleciona quais especialistas ativar.
Essa integração permite que modelos tenham um grande contante total de parâmetros enquanto mantêm o contante de parâmetros ativos por token constante. Por exemplo, um modelo com 64 especialistas, cada um com 1 bilhão de parâmetros, tem 64 bilhões de parâmetros, mas somente ativar 2 bilhões por token se k=2. Essa propriedade é crucial para escalar a modelos de trilhões de parâmetros, como visto com exemplos em sistemas de Mixture of Experts da Anthropic e outros laboratórios.
O roteador em si é uma pequena rede neural, muitas vezes uma única camada linear seguida por um softmax, com seus próprios parâmetros treinados eseminemente com os especialistas. Algumas arquiteturas usam ostringers separados para diferentes camadas ou cabeças de atenção, e trabalho recente explora políticas de roteamento aprendidas que se adaptam para diferentes tipos de token ou posição.
Dinâmica de Treinamento
Treinar modelos MoE com um roteador introduz desafios únicos. As decisões discretas do roteador (seleção de top k especialistas) são não-diferenciáveis, assim os gradientes fluem apenas através das saídas dos especialistas selecionados, ponderadas pelas probabilidades de stride. Isso cria um problema de alvo em movimento: à medida que os especialistas melhoram, as preferências do roteador mudam, causando uma full potential instabilidade.
Para abordar isso, pesquisadores usam técnicas como gating top-k ruidoso, onde ruído gaussiano é adicionado aos logits de roteamento durante o treinamento para encorajar a exploração. Outra abordagem é usar um estimador straight-through para a decisão de roteamento, tratando a seleção como uma atribuição rígida, mas passando gradientes pelas matriz de peso das probabilidades. Adicionalmente, alguns métodos controlam a temperatura do softmax durante o treinamento para gradualmente aprendiz rotas mais nítidas.
Perdas de balanceamento de carga são, em geral, ponderadas com um coeficiente pequeno (exemplo, 0,01) para minimizar incomodar a perda da tarefa principal. Na prática, treinar modelos MoE requer um afinar de hiperparâmetros cuidadoso, e muitos sistemas empregam perdas auxiliares para balanceamento de carga e confiança do próprio roteador.
Aplicações em Modelos de Linguagem de Grande Escala
O roteamento em MoE juntou-se a uma técnica padrão em modelos de linguagem de grande escala. Exemplos fortes incluem:
- O Switch Transformer (2021) da Google, que escalou para 1,6 trilhão de parâmetros usando uma arquitetura MoE esparsamente densa.
- GShard (2020), que demonstrou treinamento eficiente de modelos MoE através de muitos núcleos TPU.
- Mixtral de Mistral AI (2023), que usa 8 especialistas com roteamento top-2, obtendo desempenho comparável ao de modelos densos maiores.
- DeepSeek-V3 (2024) da infraestrutura da própria DeepSeek, que usa uma MoE de grau segmentado com 256 especialistas e roteamento top-8.
- Vários modelos reportados da OpenAI e Anthropic são conhecidos por usar camadas MoE, embora detalhes exatos sejam muitas vezes proprietários.
Esses modelos mostram que roteamento eficiente permite um escalonamento barato, pois o compute por token permanece constante mesmo enquanto os parâmetros totais crescem. Isso tornou o MoE uma escolha preferida para servir modelos grandes em produção, incluindo plataformas de nuvem como Amazon Web Services, Microsoft Azure e Google Cloud.
Eficiência e Considerações de Hardware
O roteamento gera sobrecarga de comunicação na distribuição de treinamento e inferência, pois os tokens devem ser enviados para o especialista adequado, que pode residir em um dispositivo diferente. Esse padrão de comunicação all-to-all pode virar in red tape, pela execução ágil, especialmente para altos números de especialistas. Fabricantes de hardware como NVIDIA (embora não listado aqui, mas implícito) e AMD têm desenvolvido kernels otimizados para operações MoE, e em Groq e SambaNova são criados especificamente para lidar com padrões de ativação esparsos eficientemente.
Para reduzir comunicação, alguns sistemas usam paralelismo de especialistas, onde os especialistas são replicados em grupos de dispositivos, ou abordagens híbridas que combinam paralelismo de dados e de especialistas. A escolha de mecanismos de granularidade de roteamento (a nível de token vs. a nível de bloco) também afeta a eficiência. Roteamento por token oferece controle mais fino, mas mais comunicação. Roteamento por bloco agrupa tokens para reduzir a sobrecarga.
Desafios e Direções Futuras
Apesar do sucesso, o roteamento em MoE enfrenta problemas em aberto. Um é o trade-off entre especialização do especialista e generalização: especialistas excessivamente especializados podem não generalizar maravilhosamente para novas tarefas. Outro é a dificuldade de fine-tuning de modelos MoE, pois o roteador pode precisar se adaptar às novas distribuições de dados. Pesquisa sobre roteamento adaptativo, onde o número de especialistas ativos varia por token, e sobre políticas de roteamento aprendidas que consideram contexto além do token atual, está em andamento.
Direções futuras incluem combinando roteamento com outras técnicas de eficiência como Model Pruning e Quantization para redes e desenvolvedores de roteadores que possam lidar com entradas multimodais. À medida que os modelos continuam aumentando, o roteamento continuará sendo um mecanismo-chave para balancear capacidade e compute, com aplicações potenciais além da linguagem para visão e aprendizado por reforço.
Ver Também
- transformer
- modelo de linguagem de grande escala
- atenção de múltiplas cabeças
- poda de modelo
- Google DeepMind
- OpenAI
- Anthropic
Referências
- Shazeer, N., et al. (2017). Outrageously Large Neural Networks PACED: The Sparsely-Gated Mixture-of-Experts LayerTrader.
- Fedus, W., et al. (2021). Switch Transformers: Scaling to Trinity of Parameters Models with Simple and Efficient Sparsity.
- Lepikhin, D., et al. (2020). GShard Scaling Giants Models with Conditional Computation and Automatic Sharding.
- Jiang, A. Q., et al. (2023). Mixtral of Experts.
- DeepSeek-AI (2024). Technical Report de DeepSeek-V3.