Esquema adaptativo de balanceamento de carga baseado em peso especialista para roteamento dinâmico de MoE

O desequilíbrio de carga é um grande gargalo de desempenho no treinamento de modelos mixture-of-experts (MoE), pois cargas desequilibradas entre especialistas podem levar ao colapso do roteamento. A maioria das abordagens existentes resolve esse problema introduzindo funções de perda auxiliares para balancear a carga; contudo, os hiperparâmetros dessas funções frequentemente precisam ser ajustados para diferentes tarefas. Além disso, aumentar o número de especialistas ativados tende a agravar o desequilíbrio de carga, enquanto fixar a contagem de ativação pode diminuir a confiança do modelo em lidar com tarefas difíceis. Para enfrentar esses desafios, este artigo propõe uma estratégia de roteamento balanceada de forma dinâmica que emprega um algoritmo de roteamento dinâmico baseado em limiar. Após cada etapa de roteamento, o método ajusta os pesos dos especialistas para influenciar a distribuição da carga no roteamento subsequente. Diferentemente dos métodos baseados em funções de perda para balanceamento, nossa abordagem opera diretamente no nível do roteamento, evitando perturbações de gradiente que poderiam degradar a qualidade do modelo, enquanto roteia dinamicamente para fazer uso mais eficiente dos recursos computacionais. Experimentos em benchmarks de Compreensão de Linguagem Natural (NLU) demonstram que o método proposto alcança precisão comparável ao roteamento top-2, enquanto reduz significativamente o desvio padrão da carga (por exemplo, de 12,25 para 1,18 em MNLI). Além disso, a ativação dinâmica de especialistas baseada em limiar reduz os parâmetros do modelo e oferece uma nova perspectiva para mitigar o desequilíbrio de carga entre especialistas.