UniPool: Um Pool Global de Especialistas Compartilhado para Mistura de Especialistas

Resumo

As arquiteturas modernas de Mistura de Especialistas (MoE) alocam capacidade de especialistas por meio de uma regra rígida por camada: cada camada do transformer possui um conjunto de especialistas separado. Esta convenção acopla a escalagem de profundidade com o crescimento linear de parâmetros de especialistas e assume que cada camada precisa de capacidade de especialista isolada. No entanto, análises recentes e nossa sonda de roteamento desafiam esta regra de alocação: substituir o roteador top-k aprendido de uma camada mais profunda por um roteamento aleatório uniforme reduz a precisão downstream em apenas 1,0-1,6 pontos em vários modelos MoE de produção. Motivados por esta redundância, propomos o UniPool, uma arquitetura MoE que trata a capacidade de especialistas como um orçamento arquitetônico global, substituindo a propriedade de especialistas por camada por um único pool compartilhado acessado por roteadores independentes por camada. Para permitir um treinamento estável e equilibrado sob compartilhamento, introduzimos uma perda auxiliar em nível de pool que equilibra a utilização de especialistas em todo o pool e adotamos o NormRouter para fornecer um roteamento esparso e estável em escala para o pool compartilhado de especialistas. Em cinco escalas de modelo de arquitetura LLaMA (182M, 469M, 650M, 830M e 978M de parâmetros) treinadas em 30B de tokens do Pile, o UniPool melhora consistentemente a perda de validação e a perplexidade em relação às linhas de base MoE padrão equivalentes. Nessas escalas, o UniPool reduz a perda de validação em até 0,0386 em relação ao MoE padrão. Além da melhoria bruta na perda, nossos resultados identificam o tamanho do pool como um hiperparâmetro explícito de escalagem de profundidade: variantes do UniPool com pool reduzido, usando apenas 41,6%-66,7% do orçamento de parâmetros de especialistas do MoE padrão, igualam ou superam o MoE por camada nas escalas testadas. Isso mostra que, sob um design de pool compartilhado, os parâmetros de especialistas não precisam crescer linearmente com a profundidade; eles podem crescer de forma sublinear enquanto permanecem mais eficientes e eficazes do que o MoE padrão. Uma análise adicional mostra que os benefícios do UniPool se compõem com uma decomposição de especialistas mais refinada.

English

Modern Mixture-of-Experts (MoE) architectures allocate expert capacity through a rigid per-layer rule: each transformer layer owns a separate expert set. This convention couples depth scaling with linear expert-parameter growth and assumes that every layer needs isolated expert capacity. However, recent analyses and our routing probe challenge this allocation rule: replacing a deeper layer's learned top-k router with uniform random routing drops downstream accuracy by only 1.0-1.6 points across multiple production MoE models. Motivated by this redundancy, we propose UniPool, an MoE architecture that treats expert capacity as a global architectural budget by replacing per-layer expert ownership with a single shared pool accessed by independent per-layer routers. To enable stable and balanced training under sharing, we introduce a pool-level auxiliary loss that balances expert utilization across the entire pool, and adopt NormRouter to provide sparse and scale-stable routing into the shared expert pool. Across five LLaMA-architecture model scales (182M, 469M, 650M, 830M, and 978M parameters) trained on 30B tokens from the Pile, UniPool consistently improves validation loss and perplexity over the matched vanilla MoE baselines. Across these scales, UniPool reduces validation loss by up to 0.0386 relative to vanilla MoE. Beyond raw loss improvement, our results identify pool size as an explicit depth-scaling hyperparameter: reduced-pool UniPool variants using only 41.6%-66.7% of the vanilla expert-parameter budget match or outperform layer-wise MoE at the tested scales. This shows that, under a shared-pool design, expert parameters need not grow linearly with depth; they can grow sublinearly while remaining more efficient and effective than vanilla MoE. Further analysis shows that UniPool's benefits compose with finer-grained expert decomposition.