Resumo Técnico: SlimQwen: Explorando a Poda e a Distilação no Pré-treinamento de Grandes Modelos MoE

Declaração do Problema

As arquiteturas de Mistura de Especialistas (MoE) tornaram-se dominantes para a escalabilidade de Grandes Modelos de Linguagem (LLMs), mas continuam sendo computacionalmente caras para pré-treinar e servir. Embora a poda estruturada e a distilação de conhecimento (KD) sejam técnicas estabelecidas para comprimir modelos densos, sua aplicação a modelos MoE na escala de pré-treinamento permanece incerta. Especificamente, há uma falta de compreensão sistemática sobre:

1. Se podar um MoE pré-treinado fornece uma inicialização superior em comparação com treinar uma arquitetura-alvo a partir do zero.
2. Como diferentes estratégias de compressão de especialistas (poda vs. fusão) afetam o desempenho final após o pré-treinamento contínuo em grande escala.
3. O que constitui a receita de treinamento ótima (funções de perda e agendamento) para recuperar o desempenho pós-compressão.

Metodologia

O artigo propõe o SlimQwen, um framework sistemático para comprimir modelos MoE durante o pré-treinamento. A metodologia abrange três dimensões principais:

1. Poda Estruturada e Inicialização

Os autores investigam a compressão do modelo Qwen3-Next-80A3B (um MoE de atenção híbrida) em uma configuração menor 23A2B. A compressão opera em três dimensões:

• Poda de Profundidade: Remoção das últimas 25% das camadas do transformador.
• Poda de Largura: Redução da dimensão oculta (de 2048 para 1536) com base nas estatísticas de ativação (Ativação Média Absoluta) calculadas em um conjunto de dados de calibração.
• Compressão de Especialistas: Redução do número de especialistas por módulo MoE (de 512 para 256) e do número de especialistas ativados (de 10 para 8 roteados + 1 compartilhado).

2. Estratégia de Fusão de Especialistas

Para abordar o trade-off entre preservar conhecimento especializado e consolidar especialistas descartados, os autores propõem uma estratégia de Fusão Parcial de Especialistas com Preservação. Em vez de podar puramente ou fundir agressivamente todos os especialistas:

• Os 50% superiores dos especialistas (classificados por pontuações de importância derivadas da frequência de roteamento, logits suaves ou REAP) são mantidos intactos.
• Os 50% restantes dos especialistas-alvo são construídos fundindo especialistas descartados em "especialistas base" selecionados do conjunto mantido, ponderados por suas pontuações de importância relativas.
• Esta abordagem visa prevenir a homogeneização de representações, mantendo a capacidade de absorver conhecimento de especialistas podados.

3. Receitas de Treinamento e Distilação

O artigo avalia vários objetivos de treinamento e estratégias de agendamento para o modelo comprimido:

• Funções de Perda: Os autores comparam a Distilação de Conhecimento (KD) pura contra uma abordagem híbrida que combina KD com a perda padrão de Modelagem de Linguagem (LM). Eles introduzem ainda a Distilação de Previsão de Múltiplos Tokens (MTP), que estende o objetivo de distilação para prever múltiplos tokens futuros simultaneamente, não apenas o próximo token.
• Poda Progressiva: Em vez de compressão "de uma só vez", os autores exploram agendamentos progressivos (profundidade primeiro, largura primeiro e conjunta) onde a redução estrutural ocorre em etapas (por exemplo, 40B de tokens em um estágio intermediário, seguido por 360B de tokens na configuração final-alvo).

Principais Descobertas e Resultados

Superioridade da Inicialização

Treinar uma arquitetura-alvo a partir do zero (Inicialização Aleatória) apresenta desempenho significativamente inferior em comparação com a inicialização a partir de um modelo MoE podado. Sob o mesmo orçamento de treinamento de 120B de tokens, a inicialização podada alcançou uma pontuação média de benchmark de 73,45, comparada a 61,66 para a inicialização aleatória. O modelo podado recuperou 86,5% do desempenho do professor, apesar de ser 3,4× menor.

Equivalência na Compressão de Especialistas e Preservação Parcial

• Métodos de Uma Só Vez: Após o pré-treinamento contínuo em grande escala (400B de tokens), diferentes métodos de poda ou fusão de especialistas de uma só vez (baseados em frequência, logits suaves ou REAP) convergiram para níveis de desempenho final semelhantes. Nenhuma métrica única dominou em todas as tarefas.
• Preservação Parcial: A estratégia proposta de preservação parcial (manter metade dos especialistas intactos) superou consistentemente outras estratégias de fusão em principais benchmarks (MMLU, MMLU-Pro, GSM8K), sugerindo que preservar um subconjunto de especialistas intactos ajuda a manter o conhecimento especializado.

Objetivos de Treinamento

• Perda Híbrida: Combinar a KD de Previsão do Próximo Token (NTP) com a perda padrão de LM superou a KD pura, particularmente em tarefas intensivas em conhecimento (por exemplo, MMLU melhorou de 74,16 para 74,93).
• Distilação MTP: Incorporar a KD MTP produziu ganhos consistentes em todos os benchmarks. Além disso, a KD MTP melhorou significativamente a taxa de aceitação de múltiplos tokens na decodificação especulativa (por exemplo, aumentando a aceitação de acc 1 para acc 4), aprimorando a eficiência de inferência.

Compressão Progressiva vs. de Uma Só Vez

Agendamentos de poda progressiva superaram consistentemente a compressão de uma só vez sob o mesmo orçamento total de tokens (400B). A estratégia progressiva Profundidade Primeiro (SlimQwen) alcançou os melhores resultados, melhorando as pontuações MMLU de 75,86 (uma só vez) para 77,39. Isso sugere que transições arquiteturais graduais facilitam trajetórias de otimização mais suaves para a transferência de conhecimento.

Significado e Contribuições

O artigo afirma oferecer orientações práticas para a compressão eficiente de MoE na escala de pré-treinamento. Suas principais contribuições são:

1. Estudo Sistemático da Compressão MoE: Estabelece que a poda estruturada fornece uma inicialização forte para o pré-treinamento contínuo e demonstra que diferentes métodos de compressão de especialistas de uma só vez produzem resultados semelhantes após treinamento extensivo, desde que uma estratégia adequada de recuperação seja utilizada.
2. Estratégia de Preservação Parcial: Introduz uma técnica de fusão de especialistas simples, mas eficaz, que equilibra a preservação e a consolidação de conhecimento, levando a melhorias consistentes no desempenho downstream.
3. Técnicas Avançadas de Distilação: Valida a eficácia de combinar perda LM com KD e propõe a Distilação MTP, que melhora tanto a qualidade do treinamento do backbone quanto a eficiência da decodificação especulativa.
4. Agendamento Progressivo: Demonstra que agendamentos de poda progressiva (especificamente Profundidade Primeiro) são superiores à compressão de uma só vez, confirmando que a redução de capacidade em etapas otimiza a transferência de conhecimento pré-treinado.

Empiricamente, o trabalho comprime com sucesso o modelo Qwen3-Next-80A3B em um modelo 23A2B (compressão de aprox. 4×) que mantém desempenho competitivo em benchmarks gerais de raciocínio, matemática, codificação e língua chinesa.