O roteamento esparso de MoE é um mecanismo em redes neurais de Mixture of experts que determina quais sub-redes especialistas processam cada token de entrada. Em uma camada esparsa de mistura de especialistas, um roteador ou função de gating avalia a entrada e seleciona apenas um pequeno número de especialistas, tipicamente um ou dois, de um pool maior. Isso contrasta com a mistura densa de especialistas, onde todos os especialistas contribuem para cada saída. A decisão de roteamento é feita por token, permitindo que diferentes partes de uma sequência de entrada sejam processadas por diferentes especialistas. Esse design aumenta a contagem total de parâmetros de um modelo enquanto mantém o custo computacional por token aproximadamente constante, porque apenas os especialistas selecionados são ativados durante as passagens para frente e para trás.
O conceito emergiu de trabalhos anteriores sobre mistura de especialistas na década de 1990, mas o roteamento esparso tornou-se praticamente importante nas décadas de 2010 e 2020 como uma forma de escalar grandes modelos de linguagem além dos limites das arquiteturas densas de Transformer (architecture). Ao rotear tokens para especialistas especializados, os modelos podem capturar padrões diversos nos dados sem exigir que todos os parâmetros estejam ativos para cada entrada. O roteamento esparso de MoE é agora um componente central de vários modelos de produção proeminentes, incluindo aqueles desenvolvidos por Google DeepMind, OpenAI e outras organizações de pesquisa.
Origens Históricas
A ideia fundamental da mistura de especialistas data do início da década de 1990. Pesquisadores incluindo Robert A. Jacobs, Michael I. Jordan, Steven J. Nowlan e Geoffrey E. Hinton introduziram a arquitetura de misturas adaptativas de especialistas locais em 1991, onde uma rede de gating ponderava as saídas de várias redes especialistas. Na mesma época, a rede meta-pi de John Hampshire e Alex Waibel aplicou um esquema de ponderação semelhante à classificação de fonemas de fala, treinando seis redes neurais com atraso temporal em dados de seis falantes japoneses. Esses sistemas iniciais usavam ponderação densa, significando que todos os especialistas contribuíam para cada saída, mas estabeleceram os componentes centrais: funções especialistas, uma função de gating e uma combinação ponderada.
O roteamento esparso, onde apenas um subconjunto de especialistas é ativado, foi explorado no final da década de 1990 e início da década de 2000, mas ganhou tração com o avanço do aprendizado profundo. Em 2017, pesquisadores da Google DeepMind publicaram a camada de mistura de especialistas com gating esparso, que introduziu uma rede de gating treinável que selecionava os top-k especialistas para cada entrada. Esse trabalho demonstrou que a ativação esparsa poderia aumentar dramaticamente a capacidade do modelo sem um aumento proporcional na computação, pavimentando o caminho para aplicações em larga escala posteriores.
Mecanismo de Roteamento
Em uma camada esparsa de MoE, o mecanismo de roteamento consiste em uma função de gating, tipicamente uma camada linear seguida por um softmax sobre a dimensão dos especialistas. Para uma representação de token de entrada \(x\), o roteador calcula uma pontuação para cada um dos \(n\) especialistas, frequentemente como \(w(x)_i = \text{softmax}(W_g x)_i\), onde \(W_g\) é uma matriz de pesos aprendível. O roteador então seleciona os top-k especialistas com as pontuações mais altas, onde k é um hiperparâmetro, geralmente definido como 1 ou 2. A saída da camada é uma soma ponderada das saídas dos especialistas selecionados, com pesos normalizados entre os especialistas escolhidos.
Essa seleção por token permite que o modelo se especialize: diferentes especialistas podem aprender a lidar com diferentes tipos de entradas, como diferentes idiomas, domínios ou padrões sintáticos. A decisão de roteamento é determinística durante a inferência, mas durante o treinamento, o roteador deve ser treinado em conjunto com os especialistas. Como a seleção top-k não é diferenciável, os pesquisadores usam técnicas como o estimador de passagem direta ou o gating top-k ruidoso, que adiciona ruído aprendível aos logits durante o treinamento para incentivar a exploração e o balanceamento de carga.
Balanceamento de Carga e Perdas Auxiliares
Um grande desafio no roteamento esparso de MoE é garantir que os especialistas sejam usados de forma aproximadamente igual. Sem intervenção, o roteador pode convergir para um estado onde alguns especialistas recebem a maioria dos tokens, enquanto outros permanecem subutilizados, reduzindo a capacidade efetiva. Para abordar isso, os modelos incorporam perdas auxiliares de balanceamento de carga que penalizam a distribuição desequilibrada de tokens entre os especialistas. Uma abordagem comum, introduzida no Switch Transformer (2021) por pesquisadores da Google (AI), adiciona um termo de perda proporcional à fração de tokens roteados para cada especialista multiplicada pela probabilidade média do roteador para esse especialista. Isso incentiva o roteador a distribuir tokens de forma mais uniforme.
Outra técnica é a capacidade do especialista, que limita o número de tokens que cada especialista pode processar em um lote. Se um especialista atinge sua capacidade, os tokens excedentes são descartados ou roteados para uma conexão residual. Isso impede que qualquer especialista se torne um gargalo e garante computação previsível. Métodos mais recentes, como o roteamento sem perda auxiliar usado no DeepSeek-V3 (2024), usam ajustes dinâmicos de viés para balancear a carga sem perdas auxiliares explícitas.
Escalabilidade e Eficiência
O roteamento esparso de MoE permite escalar os parâmetros do modelo muito além do que os modelos densos permitem. Por exemplo, o Switch Transformer introduziu modelos com até 1,6 trilhão de parâmetros, mas cada token ativava apenas uma pequena fração deles. Similarmente, o Mixtral 8x7B, lançado pela Mistral AI em 2023, usa oito especialistas por camada com dois ativos por token, alcançando desempenho comparável a modelos densos maiores enquanto usa menos computação por token. O GShard da Google (2020) aplicou MoE esparso à tradução automática, e modelos posteriores como GLaM (2021) e PaLM (2022) incorporaram camadas MoE para eficiência.
Os ganhos de eficiência vêm do fato de que o custo computacional por token escala com o número de especialistas ativos, não com a contagem total de parâmetros. Isso permite que os modelos tenham mais parâmetros, o que pode melhorar a precisão, enquanto mantém os custos de inferência e treinamento gerenciáveis. No entanto, modelos MoE esparsos exigem mais memória para armazenar todos os parâmetros dos especialistas, e introduzem sobrecarga de comunicação no treinamento distribuído, pois os tokens devem ser roteados para especialistas que podem residir em dispositivos diferentes.
Aplicações em Grandes Modelos de Linguagem
O roteamento esparso de MoE tornou-se uma técnica padrão em grandes modelos de linguagem. Muitos modelos de pesos abertos e proprietários usam camadas MoE para equilibrar qualidade e custo. Por exemplo, Mixtral 8x7B e Mixtral 8x22B da Mistral AI, Qwen1.5-MoE da Alibaba, e DeepSeek-V2 e V3 da DeepSeek todos empregam roteamento esparso. No espaço proprietário, o GPT-4 da OpenAI é amplamente relatado como usando uma arquitetura de mistura de especialistas, embora a empresa não tenha confirmado detalhes. Os modelos Claude da Anthropic e os modelos Gemini da Google também provavelmente incorporam camadas MoE, com base em declarações públicas e patentes.
Esses modelos usam roteamento para especializar especialistas em domínios, idiomas ou tarefas de raciocínio. Por exemplo, alguns especialistas podem lidar com raciocínio matemático, enquanto outros focam em geração de código ou texto multilíngue. O roteador aprende a direcionar tokens para especialistas apropriados com base no contexto de entrada, melhorando o desempenho geral sem exigir que todos os especialistas processem cada token.
Desafios de Treinamento
Treinar modelos MoE esparsos apresenta desafios únicos além do balanceamento de carga. O roteador pode sofrer de instabilidade, onde pequenas mudanças nos parâmetros levam a grandes mudanças nas decisões de roteamento, causando divergência no treinamento. Técnicas como gating top-k ruidoso e dropout no roteador ajudam a estabilizar o treinamento. Além disso, os especialistas podem se especializar cedo demais, levando a uma forma de colapso de modo onde o modelo falha em generalizar. Pesquisadores propuseram métodos como dropout de especialistas e regularização de roteamento para mitigar isso.
Outro desafio é a sobrecarga de comunicação no treinamento distribuído. Em uma configuração típica, cada especialista é colocado em um dispositivo diferente, e os tokens devem ser enviados de seu dispositivo original para o dispositivo que hospeda o especialista selecionado. Essa comunicação all-to-all pode se tornar um gargalo, especialmente para tamanhos de lote grandes. Implementações eficientes, como aquelas nas bibliotecas Megatron e DeepSpeed, usam cronogramas de comunicação otimizados para reduzir essa sobrecarga.
Desenvolvimentos Recentes
Pesquisas recentes focaram em melhorar a qualidade do roteamento e reduzir o custo dos modelos MoE. A abordagem Mixture-of-Depths (2024) estende a ideia de esparsidade à dimensão de profundidade, roteando tokens não apenas para diferentes especialistas, mas também pulando camadas inteiras. Outra direção é a segmentação fina de especialistas, onde os especialistas são menores e mais numerosos, permitindo especialização em granularidade mais fina. O DeepSeek-V3, por exemplo, usa 256 especialistas por camada com apenas 8 ativos, alcançando desempenho de ponta a custo de treinamento reduzido.
O co-design de hardware e software também avançou. Empresas como NVIDIA e AMD otimizaram seus aceleradores para inferência MoE, e frameworks como vLLM e TensorRT-LLM suportam serviço MoE eficiente. O crescente interesse em implantação de borda levou a pesquisas sobre compressão de modelos MoE, como poda de especialistas e quantização, para caber em dispositivos com memória limitada.
Direções Futuras
O roteamento esparso de MoE permanece uma área ativa de pesquisa. Perguntas em aberto incluem como determinar automaticamente o número ideal de especialistas e a granularidade do roteamento, como tornar o roteamento mais interpretável, e como combinar MoE com outras técnicas de eficiência como poda de modelos e quantização. À medida que os modelos continuam a escalar, o roteamento esparso provavelmente desempenhará um papel central no equilíbrio entre capacidade e custo computacional, permitindo sistemas de IA mais poderosos sem demandas de recursos proibitivas.