ConMoE: Expert-Pool Consolidation via Prototype Reassignment for MoE Compression

Resumo Técnico do ConMoE: Consolidação do Conjunto de Especialistas via Reatribuição de Protótipos para Compressão de MoE

1. Declaração do Problema

Modelos de linguagem do tipo Mixture-of-Experts (MoE) alcançam escalabilidade ativando apenas um pequeno subconjunto de especialistas por token, reduzindo o cálculo por token. No entanto, o conjunto completo de especialistas roteados ainda deve ser armazenado e servido, criando um gargalo significativo de memória para implantação. Os métodos existentes de compressão pós-treinamento abordam isso principalmente através de poda de especialistas (remoção de especialistas com base no uso ou importância) ou fusão de especialistas (combinação de múltiplos especialistas em novos módulos fundidos).

Os autores argumentam que essas abordagens frequentemente confundem duas decisões distintas:

1. Quais parâmetros de especialista devem ser retidos?
2. Como as referências originais do roteador aos especialistas devem ser representadas após a compressão?

A poda remove parâmetros, enquanto a fusão constrói novos pesos, frequentemente exigindo ajuste fino pós-compressão ou operadores de fusão complexos. O artigo propõe uma visão complementar: Consolidação do Conjunto de Especialistas. Esta abordagem retém um conjunto menor de especialistas pré-treinados como protótipos reutilizáveis e remapeia deterministicamente cada referência original de especialista para um desses protótipos selecionados. Isso separa o armazenamento do conjunto reduzido da estrutura de reutilização, permitindo que a interface original do roteador permaneça inalterada sem atualizações de pesos.

2. Metodologia: ConMoE

O framework proposto, ConMoE, é um método de compressão sem treinamento e de um único passo que opera via seleção de protótipos e reatribuição determinística.

A. Formulação do Problema
O método define um escopo local $G$ contendo uma ou mais camadas MoE vizinhas. Dado uma taxa de redução $\rho$, o objetivo é construir um conjunto de protótipos reduzido $P_G$ de tamanho $K$ a partir do conjunto original de especialistas $E_G$. É estabelecido um mapa de reatribuição $m_G: E_G \to P_G$, onde cada especialista original $e$ é mapeado para um protótipo selecionado $p$.
O objetivo ideal é minimizar o custo de reatribuição $L_G(P)$, equilibrando a importância dos especialistas retidos contra o erro introduzido pelo remapeamento dos demais.

B. Pontuação e Seleção de Protótipos
Em vez de resolver diretamente o problema de otimização combinatória, o ConMoE utiliza uma regra de pontuação heurística para selecionar os $K$ melhores especialistas como protótipos. Para cada especialista $e$, duas pontuações são calculadas usando dados de calibração:

1. Contribuição Condicionada ao Roteamento ($a_e$): Mede a contribuição média de um especialista quando selecionado pelo roteador, ponderada pelas probabilidades de roteamento e normas de saída.
2. Substituibilidade ($b_e$): Mede a distância até o vizinho mais próximo dentro do escopo. Uma distância maior implica que o especialista é mais difícil de substituir.

A pontuação final é o produto das pontuações normalizadas de contribuição e substituibilidade ($s_e = \bar{a}_e \bar{b}_e$). Os $K$ melhores especialistas são selecionados como o conjunto de protótipos $P_G$.

C. Reatribuição Determinística
Uma vez selecionados os protótipos, cada especialista original é atribuído ao seu protótipo mais próximo em $P_G$.

• Sem Atualização de Pesos: Os protótipos selecionados são reutilizados diretamente como especialistas pré-treinados; não ocorre fusão de pesos ou ajuste fino no pipeline padrão.
• Preservação do Roteador: O roteador original permanece inalterado. Quando o roteador seleciona um conjunto de especialistas para um token, o operador consolidado redireciona essas chamadas para seus protótipos atribuídos. Se múltiplos especialistas roteados mapearem para o mesmo protótipo, seus pesos de roteamento são agregados em um único coeficiente.

D. Escopo e Reutilização Inter-camadas
O ConMoE suporta compartilhamento local inter-camadas. Um escopo pode conter múltiplas camadas vizinhas, permitindo que protótipos sejam compartilhados entre elas se exibirem redundância. No entanto, o escopo é limitado para evitar o descompasso de transformações específicas de profundidade encontradas em camadas distantes.

3. Principais Contribuições

1. Formulação: O artigo formula a compressão de MoE de um único passo como consolidação do conjunto de especialistas, separando explicitamente a decisão de quais parâmetros armazenar de como as slots originais do roteador são representadas.
2. Framework (ConMoE): Um método de remapeamento sem treinamento que preserva a interface original do roteador enquanto reduz o conjunto lógico de especialistas roteados. Baseia-se em sinais de contribuição e substituibilidade baseados em calibração para selecionar protótipos sem atualizações de gradiente.
3. Validação Empírica: Os autores demonstram que a consolidação baseada em remapeamento é uma alternativa viável à poda e à fusão em múltiplos modelos MoE pré-treinados (Qwen3, DeepSeekMoE, OLMoE).
4. Insights de Ablação:
   • A reatribuição determinística é identificada como o componente mais estável do framework.
   • O compartilhamento inter-camadas é benéfico, mas estritamente local; expandir escopos excessivamente (por exemplo, 8–16 camadas) degrada o desempenho devido ao descompasso de profundidade.
   • A fusão de pesos a posteriori (por exemplo, Arcee, média ponderada) mostrou-se menos estável do que a reutilização direta de protótipos neste cenário.

4. Resultados Experimentais

Os experimentos foram conduzidos em três modelos: Qwen3-30B-A3B, deepseek-moe-16b-base e OLMoE-1B-7B-0125. A compressão foi avaliada em taxas de redução lógica de especialistas roteados de 25% e 50%, utilizando seis benchmarks de múltipla escolha (WinoGrande, ARC-C, ARC-E, BoolQ, HellaSwag, PIQA).

• Desempenho: O ConMoE iguala ou supera bases fortes de poda (Frequency, REAP) e fusão (M-SMoE, HC-SMoE).
  • No DeepSeekMoE, o ConMoE alcançou a melhor pontuação média em ambas as reduções de 25% e 50%.
  • No OLMoE, performou melhor na redução de 25% e permaneceu competitivo em 50%.
  • No Qwen3, embora a poda fosse forte (especialmente em 50%), o ConMoE permaneceu competitivo e preservou desempenho robusto em tarefas específicas como WinoGrande e BoolQ.
• Ablações:
  • Tamanho do Escopo: Escopos pequenos (1–4 camadas) produziram resultados comparáveis ou melhores. Escopos maiores (8–16 camadas) causaram quedas significativas de desempenho, confirmando que a substituibilidade de especialistas é local.
  • Política de Seleção: A contribuição condicionada ao roteamento foi o sinal dominante. A seleção baseada apenas em distância performou mal, e a substituibilidade serviu melhor como regularizador para alocação de capacidade.
  • Fusão: O remapeamento direto de protótipos superou métodos que aplicavam fusão de pesos a posteriori, sugerindo que os ganhos provêm da própria estrutura de remapeamento e não da fusão de pesos.

5. Significado e Alegações

O artigo alega que a reatribuição determinística de protótipos oferece uma alternativa robusta e sem treinamento à poda e fusão tradicionais. Ao tratar a compressão como uma consolidação do conjunto de especialistas em protótipos reutilizáveis, o ConMoE alcança economias significativas de memória (redução lógica) enquanto mantém a interface do roteador do modelo original e evita a instabilidade frequentemente associada à fusão de pesos ou ajuste fino.

Os autores enfatizam que sua abordagem é modesta quanto ao seu escopo:

• É sem treinamento, mas depende de dados de calibração para estatísticas de roteamento.
• A reutilização inter-camadas é local; não assume intercambiabilidade global de especialistas em toda a profundidade do modelo.
• As economias relatadas são lógicas; realizar economias físicas de memória na implantação requer um runtime de protótipo compartilhado ou formato de checkpoint, e não apenas um checkpoint materializado.

O trabalho distingue-se ao mostrar que a "estrutura de reutilização" pode ser uma parte explícita do modelo comprimido, em vez de um subproduto de exclusão ou fusão, fornecendo uma base estável para compressão de MoE.