Estratégia de Balanceamento de Carga sem Perda Auxiliar para Mistura de Especialistas

Para modelos de Mistura de Especialistas (MoE), uma carga desequilibrada entre especialistas leva ao colapso do roteamento ou ao aumento do custo computacional. Métodos existentes geralmente empregam uma perda auxiliar para incentivar o equilíbrio de carga, mas uma perda auxiliar grande introduz gradientes de interferência significativos durante o treinamento, comprometendo o desempenho do modelo. Para controlar o equilíbrio de carga sem gerar gradientes indesejados durante o treinamento, propomos o Balanceamento sem Perda, caracterizado por uma estratégia de balanceamento de carga sem perda auxiliar. Especificamente, antes da decisão de roteamento top-K, o Balanceamento sem Perda aplica um viés por especialista aos scores de roteamento de cada especialista. Atualizando dinamicamente o viés de cada especialista conforme sua carga recente, o Balanceamento sem Perda mantém consistentemente uma distribuição equilibrada de carga entre especialistas. Além disso, como não gera gradientes de interferência, ele também eleva o limite superior do desempenho do modelo obtido com o treinamento MoE. Validamos o desempenho do Balanceamento sem Perda em modelos MoE com até 3B parâmetros treinados em até 200B tokens. Resultados experimentais mostram que o Balanceamento sem Perda alcança melhor desempenho e melhor equilíbrio de carga em comparação com estratégias tradicionais de balanceamento controladas por perda auxiliar.