UniPool reduz orçamento de parâmetros MoE em 34 a 58 por cento

Uma equipe de oito pesquisadores publicou UniPool, uma arquitetura MoE que substitui silos de especialistas por camada por um único pool de especialistas compartilhado globalmente. O resultado principal: variantes UniPool com pool reduzido correspondem ou superam modelos MoE full layer-wise enquanto usam apenas 41,6–66,7% do orçamento de parâmetros de especialistas vanilla.

O problema central que UniPool aborda é uma suposição incorporada a cada design MoE importante: cada camada transformer precisa de seu próprio conjunto isolado de especialistas. Esse acoplamento força a contagem de parâmetros de especialistas a escalar linearmente com a profundidade. Para testar se a suposição se sustenta, os autores substituíram roteadores top-k treinados em camadas mais profundas por roteamento uniforme aleatório e mediram a precisão em modelos MoE de produção. A queda foi apenas 1,0–1,6 pontos em múltiplos modelos, confirmando que a alocação de especialistas em camadas mais profundas é amplamente redundante.

A implementação do UniPool substitui a propriedade de especialistas por camada por um único pool consultado por roteadores independentes por camada. Dois mecanismos acompanham o design: uma perda auxiliar no nível do pool que equilibra a utilização em todo o pool em vez de dentro de camadas individuais, e NormRouter, que fornece sinais de roteamento esparsos e estáveis em escala para o pool compartilhado. A perda auxiliar impede que especialistas individuais monopolizem o tráfego conforme o pool é consultado de cada profundidade.

A equipe validou UniPool em cinco escalas de arquitetura LLaMA—182M, 469M, 650M, 830M e 978M parâmetros—cada uma treinada em 30 bilhões de tokens do Pile. UniPool melhora a perda de validação e a perplexidade sobre baselines MoE vanilla correspondentes em cada escala. A redução máxima de perda de validação em relação ao MoE vanilla é 0,0386, alcançada na maior escala testada.

Para equipes de infraestrutura de IA corporativa, as implicações pousam principalmente na camada de model-serving. Arquiteturas MoE atraem adoção porque as contagens de parâmetros ativados são baixas em relação ao tamanho total do modelo, mas a contagem total de parâmetros ainda impulsiona a pegada de memória e o tamanho do checkpoint. Uma lei de crescimento de especialistas sublinear reduz ambas. O achado também afeta a economia de fine-tuning: atualizar um pool compartilhado em vez de conjuntos de especialistas por camada reduz o número de matrizes de peso distintas que devem ser checkpointadas ou adaptadas via métodos estilo LoRA.

Questões abertas permanecem sobre a estabilidade de roteamento de produção em escalas além de 1B parâmetros, e a avaliação baseada em Pile do paper antecede convenções recentes de domain-mix. Se a estabilidade de escala do NormRouter se sustenta sob distribuições de tokens de longo contexto ou multimodal é não testado. Os autores observam que os benefícios do UniPool se compõem com técnicas de decomposição de especialistas mais refinadas, deixando espaço para designs híbridos.

O resultado de redundância de roteamento sozinho—menos de 1,6 pontos perdidos para roteamento aleatório em profundidade—dá aos designers MoE uma desculpa baseada em princípios para parar de pagar a taxa de especialistas por camada.