Resumo Técnico: Transferência de Hiperparâmetros com Camadas de Mistura de Especialistas

1. Declaração do Problema

Camadas de Mistura de Especialistas (MoE) tornaram-se um mecanismo crítico para escalar redes neurais modernas, desacoplando o número total de parâmetros treináveis dos parâmetros ativos durante a passagem forward. No entanto, o treinamento de modelos MoE esparsos introduz complexidade significativa na seleção de hiperparâmetros (HP). Diferentemente de modelos densos, arquiteturas MoE introduzem novos parâmetros treináveis (pesos do roteador) e novas dimensões arquitetônicas (número de especialistas, tamanho do especialista) que exigem ajuste cuidadoso.

Ajustar diretamente hiperparâmetros (como taxa de aprendizado, escala de inicialização e decaimento de peso) em grandes escalas é computacionalmente proibitivo. Embora técnicas de transferência de HP existam para transformers densos — permitindo que hiperparâmetros ótimos encontrados em modelos pequenos sejam extrapolados para modelos maiores —, esses métodos não foram rigorosamente adaptados para as dinâmicas de escalonamento específicas das camadas MoE esparsas. O desafio central é determinar como escalar hiperparâmetros ao aumentar simultaneamente a largura, a profundidade, a contagem de especialistas e o tamanho do especialista, sem reajustar em cada escala, garantindo ao mesmo tempo estabilidade e desempenho no treinamento.

2. Metodologia

2.1. Parametrização Proposta

Os autores propõem uma nova parametrização para modelos transformer com camadas MoE, estendendo a parametrização CompleteP (anteriormente desenvolvida para transformers densos) para incluir regras de escalonamento específicas para MoE. O objetivo é definir um conjunto de regras que prevejam como os valores brutos de hiperparâmetros (desvio padrão de inicialização $\sigma$ e taxa de aprendizado $\eta$) devem mudar conforme as dimensões do modelo escalam, de modo que as dinâmicas de treinamento permaneçam consistentes.

A derivação baseia-se no princípio da parametrização de máxima atualização ($\mu$P), que exige que os componentes da rede (pré-ativações e atualizações residuais) permaneçam $O(1)$ na inicialização e recebam atualizações de $\Theta(1)$ por etapa de treinamento. Os autores estendem isso para MoEs, exigindo que as condições de máxima atualização se mantenham não apenas para a saída da camada, mas para componentes individuais de especialistas (coeficientes de mistura e saídas dos especialistas).

Regras de escalonamento chave derivadas para o módulo MoE incluem:

• Pesos do Roteador: A taxa de aprendizado escala como $\eta \propto n_{embd}^{-1}$. A inicialização escala como $n_{embd}^{-\gamma}$ (com $\gamma \geq 0.5$).
• Vieses dos Especialistas: Inicializados em zero com uma taxa de aprendizado constante $\eta \propto 1$ (independente da contagem de especialistas), desde que a esparsidade seja fixa.
• Pesos MLP dos Especialistas:
  • Projeção de subida ($W_{up}$): $\sigma_{init} \propto n_{embd}^{-1/2}$, $\eta \propto n_{embd}^{-1}$.
  • Projeção de descida ($W_{down}$): $\sigma_{init} \propto \alpha_{ffn}^{-1} n_{embd}^{-1/2}$, $\eta \propto \alpha_{ffn}^{-1} n_{embd}^{-1}$.
  • Aqui, $\alpha_{ffn}$ representa o multiplicador do tamanho oculto do especialista em relação à dimensão de incorporação.

Crucialmente, os autores fixam a razão de esparsidade $\kappa = n_{act}/n_{exp}$ (a fração de especialistas ativos) como uma constante ao escalar o número total de especialistas ($n_{exp}$) e especialistas ativos ($n_{act}$). Isso contrasta com abordagens que fixam o número de especialistas ativos enquanto aumentam o pool total.

2.2. Fundamentação Teórica: Teoria de Campo Médio Dinâmico (DMFT)

Para justificar essas regras de escalonamento heurísticas, os autores empregam a Teoria de Campo Médio Dinâmico (DMFT). Eles analisam as dinâmicas de treinamento de redes residuais com camadas MoE no limite simultâneo de largura infinita ($n_{embd}$), profundidade ($L$), tamanho do especialista ($n_{hid}$) e contagem de especialistas ($n_{exp}$), mantendo uma esparsidade de ativação constante $\kappa$.

A análise revela uma nova hierarquia de campo médio de três níveis:

1. Representações do Fluxo Residual: Campo médio sobre as saídas dos especialistas.
2. Saídas dos Especialistas: Campo médio sobre os neurônios individuais dos especialistas.
3. Neurônios Individuais: Campo médio dentro de cada especialista.

A análise DMFT demonstra que, sob a parametrização proposta, as dinâmicas de treinamento no limite são:

• Independentes da razão FFN ($\alpha_{ffn}$): As dinâmicas não dependem do tamanho específico dos especialistas em relação à dimensão de incorporação, desde que o limite de escalonamento conjunto seja tomado.
• Dependentes apenas da esparsidade $\kappa$: As dinâmicas são consistentes em todos os parâmetros de escalonamento, desde que a razão de esparsidade seja fixa.
• Invariante à Escala: A evolução das estatísticas de resumo da rede (por exemplo, kernels de características por camada) é consistente entre escalas, garantindo teoricamente que os hiperparâmetros sejam transferidos de forma confiável.

2.3. Configuração Experimental

Os autores validam empiricamente sua parametrização usando modelos de linguagem Transformer apenas decodificador nos conjuntos de dados FineWeb e C4.

• Modelos Base: Ajustados em modelos pequenos com aproximadamente 38M de parâmetros ativados.
• Escalonamento: Escalonados para modelos com até 2B de parâmetros totais, variando largura, profundidade, contagem de especialistas e tamanho do especialista.
• Orçamento de Tokens: Os experimentos foram conduzidos em um orçamento fixo de 1B de tokens (2000 etapas) para isolar as dinâmicas do treinamento inicial, bem como horizontes mais longos (até 7,5B de tokens).
• Otimizador: Otimizador Adam padrão.
• Balanceamento de Carga: Uma estratégia sem perda auxiliar foi utilizada, atualizando vieses de especialistas diretamente para incentivar o balanceamento de carga sem adicionar um termo de regularização à função de perda.

3. Contribuições Principais

1. Parametrização MoE: O artigo estende a parametrização CompleteP para modelos MoE esparsos, fornecendo regras de escalonamento explícitas para pesos do roteador, vieses de especialistas e pesos MLP de especialistas através de largura, profundidade, contagem de especialistas e tamanho do especialista.
2. Justificativa Teórica via DMFT: Os autores fornecem uma fundamentação teórica rigorosa para sua parametrização usando DMFT. Eles derivam uma descrição explícita das dinâmicas de treinamento no limite de largura/profundidade infinita, provando que as dinâmicas convergem para um sistema estável e invariante à escala que depende apenas da esparsidade, e não do escalonamento específico das dimensões dos especialistas.
3. Validação Empírica da Transferência de HP: O estudo demonstra que hiperparâmetros ótimos (taxa de aprendizado e escala de inicialização) identificados em modelos base pequenos (38M de parâmetros ativos) são transferidos de forma confiável para modelos muito maiores (até 2B de parâmetros totais) através de várias dimensões arquitetônicas.
4. Insights Arquiteturais: Os autores verificam empiricamente que, sob sua parametrização, aumentar o número de especialistas (mantendo os parâmetros totais fixos) produz melhor desempenho do que aumentar o tamanho dos especialistas individuais. Essa descoberta alinha-se com a literatura recente, mas é alcançada aqui sem a necessidade de varreduras caras de hiperparâmetros em cada escala.

4. Resultados

• Transferência Confiável de HP: Sob as regras de escalonamento propostas, taxas de aprendizado base ótimas e desvios padrão de inicialização transferem-se efetivamente entre modelos que variam de 51M a 2B de parâmetros totais. As curvas de perda dos modelos escalados colapsam com o modelo base nas iterações iniciais antes de divergir (com modelos maiores alcançando perda menor).
• Estabilidade: A parametrização garante dinâmicas de treinamento estáveis, incluindo balanceamento de carga uniforme de especialistas, mesmo ao escalar o número de especialistas. Os autores observam que o pré-treinamento MoE é particularmente sensível a hiperparâmetros de escala constante (multiplicadores tratados como $\Theta(1)$), e o ajuste desses é crucial para a estabilidade.
• Desempenho: Modelos treinados usando hiperparâmetros zero-shot (transferidos de modelos pequenos) alcançam desempenho competitivo contra baselines densas (por exemplo, GPT-2 small/medium) quando comparados em contagens de parâmetros ativos.
• Contagem de Especialistas vs. Tamanho: Experimentos confirmam que aumentar o número de especialistas é mais eficiente em parâmetros do que aumentar o tamanho do especialista em uma contagem de parâmetros fixa. Esse benefício persiste mesmo em horizontes de treinamento mais longos (5B de tokens).
• Sensibilidade à Esparsidade: O estudo destaca que a transferência de hiperparâmetros é válida apenas quando a razão de esparsidade $\kappa$ é fixa. Escalar o número de especialistas enquanto se fixa o número de especialistas ativos (enviando $\kappa \to 0$) quebra a transferibilidade dos hiperparâmetros ótimos.

5. Significado e Alegações

O artigo alega fornecer uma estrutura prática, porém rigorosa para escalar modelos MoE. Ao combinar parametrização heurística com análise DMFT, os autores oferecem um método para:

• Reduzir Custos de Treinamento: Permitir a seleção de hiperparâmetros ótimos para modelos MoE em grande escala ajustando apenas modelos base pequenos e baratos.
• Garantir Estabilidade: Fornecer regras que previnem instabilidade no treinamento (por exemplo, colapso ou divergência de especialistas), comum no treinamento de MoE esparsos.
• Guiar o Design Arquitetural: Oferecer evidências teóricas e empíricas de que aumentar a contagem de especialistas é preferível a aumentar o tamanho do especialista sob esparsidade fixa, auxiliando no design de modelos eficientes em grande escala.

Os autores reconhecem limitações, observando que o trabalho atual foca na taxa de aprendizado e inicialização, deixando outros hiperparâmetros (tamanho de lote, decaimento de peso, cronogramas de LR) para investigação futura. Eles também notam que, embora a análise DMFT suporte o limite de largura infinita, o comportamento de transferências de largura pequena requer estudo teórico adicional. Além disso, o artigo não alega resolver as leis de escalonamento "ótimas em computação" para MoEs (por exemplo, expoentes Chinchilla), pois o compromisso FLOP-desempenho em MoEs difere significativamente de modelos densos devido a restrições de hardware induzidas pela esparsidade.

Em resumo, este trabalho estabelece uma parametrização fundamental que permite a extrapolação confiável de hiperparâmetros de modelos MoE pequenos para grandes, fundamentada em uma nova teoria de campo médio de três níveis das dinâmicas de treinamento.