ProbMoE: Differentiable Probabilistic Routing for Mixture-of-Experts

Resumo Técnico: ProbMoE: Roteamento Probabilístico Diferenciável para Mistura de Especialistas

Declaração do Problema
Modelos de Mistura de Especialistas (MoE) escalam eficientemente ao ativar apenas um pequeno subconjunto de especialistas por token. No entanto, o treinamento desses modelos permanece desafiador porque o roteamento padrão depende de um operador top-$k$ discreto e não diferenciável. Os procedimentos de treinamento convencionais contornam isso ignorando a dependência da seleção top-$k$ em relação aos logits do roteador, propagando gradientes apenas através das probabilidades softmax dos especialistas selecionados. Essa aproximação limita o sinal de aprendizado para subconjuntos de especialistas alternativos, levando frequentemente a distribuições de roteamento concentradas, utilização pobre dos especialistas e dinâmica de treinamento instável. Tentativas existentes para mitigar isso, como retropropagação densa ou estimadores de gradiente heurísticos, muitas vezes falham ao não modelar explicitamente a distribuição sobre subconjuntos de especialistas viáveis ou ao fornecer mecanismos limitados para explorar combinações alternativas durante o treinamento.

Metodologia
Os autores propõem o ProbMoE, um framework que reformula o roteamento de especialistas como um problema de inferência probabilística sobre subconjertos de especialistas com restrição de cardinalidade. Em vez de tratar o roteamento como um operador determinístico, o ProbMoE modela a seleção de especialistas como uma distribuição sobre subconjuntos discretos.

1. Formulação Probabilística:

   • Roteamento Exact-$k$: O framework define uma distribuição sobre subconjetos $S$ onde $|S|=k$. Ele condiciona as probabilidades de seleção Bernoulli independentes (derivadas dos logits do roteador via sigmoide) à restrição de que exatamente $k$ especialistas sejam selecionados. A constante de normalização para esta distribuição, $Z_k$, é computada exatamente usando programação dinâmica, tornando a formulação tratável.
   • Roteamento Dynamic-$k$: O framework se generaliza para um cenário de restrição de intervalo, onde o tamanho do subconjunto $|S|$ reside em $[k_{min}, k_{max}]$. A distribuição é condicionada a este intervalo, permitindo que o modelo infira conjuntamente as identidades dos especialistas e a cardinalidade do roteamento para cada token.

2. Passagem de Forward (Propagação Direta):

   • Durante o treinamento, o ProbMoE amostra um subconjunto de especialistas $S$ a partir da distribuição definida (exact-$k$ ou restrita por intervalo).
   • A camada MoE computa a saída usando apenas os especialistas amostrados, preservando o padrão de execução esparsa dos modelos MoE padrão.
   • Na inferência, o conjunto Máximo-A-Posteriori (MAP) é selecionado, garantindo que os custos de inferência correspondam ao roteamento MoE convencional.

3. Passagem de Backward (Propagação de Erro):

   • Para permitir a otimização baseada em gradiente, o ProbMoE emprega o estimador de gradiente SIMPLE (Ahmed et al., 2023).
   • Em vez de diferenciar através da amostra discreta, a passagem de backward utiliza as probabilidades marginais de cada especialista ser selecionado sob a distribuição restrita.
   • Os pesos de roteamento são construídos usando um estimador straight-through: $w = (\text{stopgrad}(z - m) + m) \odot \pi$, onde $z$ é a máscara amostrada, $m$ é o vetor das probabilidades marginais exatas e $\pi$ são as probabilidades de roteamento softmax.
   • Isso permite que os gradientes fluam através das probabilidades marginais, fornecendo atualizações informativas ao roteador que refletem como mudanças nas probabilidades de seleção afetam a perda esperada, enquanto mantém a execução esparsa dos especialistas.

Principais Contribuições

1. Formulação de Inferência Probabilística: O artigo trata o roteamento MoE como inferência probabilística sobre subconjuntos de especialistas com restrição de cardinalidade, indo além de aproximações heurísticas de gradiente.
2. Otimização Baseada em Marginais Tratável: Ao aproveitar o estimador SIMPLE, os autores combinam amostragem estocástica de subconjuntos com gradientes de roteador baseados em marginais exatos, permitindo o treinamento diferenciável sem computação densa do lado do especialista.
3. Framework Unificado Dynamic-$k$: O método se estende naturalmente para o roteamento dinâmico, onde o modelo infere tanto as identidades dos especialistas quanto o número de especialistas ativos por token dentro de um intervalo predefinido, sem depender de regras de gating heurísticas.
4. Melhorias Empíricas:
   • Exact-$k$: O ProbMoE alcança forte desempenho em benchmarks (GSM8K, Law, Translation, MBPP, MMLU) em backbones OLMoE e Qwen, frequentemente igualando ou superando baselines competitivos como DenseMixer e Top-$k$ Convencional.
   • Dynamic-$k$: O ProbMoE Dynamic-$k$ alcança desempenho comparável ao Exact-$k$ enquanto ativa menos especialistas em média, demonstrando alocação adaptativa eficaz.
5. Diversidade de Roteamento: A análise mostra que o ProbMoE induz distribuições de roteamento mais amplas e uma utilização de especialistas mais equilibrada em comparação com os baselines, mitigando o colapso de especialistas.

Resultados e Significância
O artigo demonstra que o ProbMoE melhora a utilização dos especialistas e a diversidade de roteamento. Em estudos de ablação, a combinação de amostragem estocástica e otimização baseada em marginais é mostrada como o principal motor dos ganhos de desempenho, superando configurações que utilizam apenas um desses componentes.

Os autores afirmam que o ProbMoE fornece um método principista para seleção de especialistas que melhora a estabilidade do treinamento e a eficiência computacional. Ao modelar o roteamento probabilisticamente, o framework permite a exploração de combinações alternativas de especialistas durante o treinamento, levando a distribuições de roteamento melhor calibradas. A variante Dynamic-$k$ mostra ainda que a cardinalidade adaptativa pode ser aprendida para corresponder à complexidade do nível do token, reduzindo os custos efetivos de roteamento sem sacrificar o desempenho da tarefa. O trabalho sugere que direções futuras incluem a extensão do ProbMoE para pré-treinamento e otimizações de nível de sistema para explorar totalmente a esparsidade em cenários dinâmicos.