MaskMoE: Potencializando o Aprendizado ao Nível de Token via Máscara de Roteamento em Mistura de Especialistas

Aumentar a capacidade do modelo aprimora suas capacidades, mas aumenta significativamente o custo computacional. Modelos de Mistura de Especialistas (MoEs) abordam isso permitindo que a capacidade do modelo escale sem aumentar substancialmente os custos de treinamento ou inferência. Apesar dos resultados promissores, modelos MoE enfrentam vários desafios. Principalmente, a dispersão dos tokens de treinamento entre múltiplos especialistas pode levar ao subajuste, particularmente para tokens infrequentes. Além disso, embora mecanismos fixos de roteamento possam mitigar esse problema, eles comprometem a diversidade das representações. Neste artigo, propomos MaskMoE, um método projetado para aprimorar o aprendizado ao nível de token utilizando uma técnica de máscara de roteamento dentro do modelo de Mistura de Especialistas. MaskMoE é capaz de manter a diversidade das representações enquanto alcança um treinamento mais abrangente. Resultados experimentais demonstram que nosso método supera modelos anteriores dominantes de Mistura de Especialistas tanto em perplexidade (PPL) quanto em tarefas downstream.