Resumo Técnico: MONA – Otimizador de Muon com Aceleração Nesterov

1. Declaração do Problema

O treinamento de modelos de linguagem de grande escala (LLMs) exige cada vez mais otimizadores com eficiência superior de amostragem. Embora o AdamW tenha sido o padrão por mais de uma década, alternativas recentes como o Muon surgiram. O Muon melhora os métodos de primeira ordem tratando matrizes de pesos como unidades geométricas únicas e aplicando ortogonalização de matrizes (via iteração Newton-Schulz) ao buffer de momento. Essa abordagem alinha-se à descida mais íngreme sob a norma espectral, oferecendo aproximadamente $2\times$ de eficiência computacional sobre o AdamW no treinamento de Misturas de Especialistas (MoE) em grande escala.

No entanto, o Muon compartilha uma limitação fundamental com todos os métodos de gradiente de primeira ordem: carece de mecanismos explícitos para explorar a paisagem de perda além da informação local do gradiente. Consequentemente, o Muon pode ficar preso em mínimos locais agudos, particularmente em cenários de treinamento com grandes lotes. Mínimos agudos estão associados a uma generalização pobre, enquanto mínimos mais planos tendem a produzir melhor desempenho em tarefas subsequentes. Embora o otimizador ALTO tenha introduzido um termo de aceleração baseado na média móvel exponencial (EMA) das diferenças de gradiente para escapar de mínimos agudos, seu design como um adaptador geral para Adam e LAMB apresenta desafios de integração com o processo de ortogonalização altamente não linear do Muon.

2. Metodologia: MONA

Os autores propõem o MONA (Otimizador de Muon com Aceleração Nesterov), que integra perfeitamente um termo de aceleração consciente da curvatura diretamente na estrutura de ortogonalização do Muon.

Algoritmo Central

O MONA modifica o pipeline padrão do Muon introduzindo um termo de aceleração antes da acumulação de momento e da ortogonalização. O processo de atualização para um parâmetro matricial $W$ na iteração $k$ procede da seguinte forma:

1. Cálculo da Diferença de Gradiente: Calcule a diferença entre o gradiente estocástico atual $G_k$ e o gradiente anterior $G_{k-1}$:
   $$D_k = G_k - G_{k-1}$$
2. Atualização do Buffer de Aceleração: Mantenha uma média móvel exponencial dessas diferenças:
   $$A_k = \beta_a A_{k-1} + (1 - \beta_a) D_k$$
   Aqui, $\beta_a$ controla a memória da aceleração.
3. Gradiente Consciente da Curvatura: Construa um gradiente aumentado $\tilde{G}_k$ adicionando o termo de aceleração escalado ao gradiente bruto:
   $$\tilde{G}_k = G_k + \alpha A_k$$
   O coeficiente $\alpha$ escala a aceleração. A motivação teórica sugere que $G_k - G_{k-1}$ aproxima $H_k(\theta_k - \theta_{k-1})$, onde $H_k$ é o Hessiano. Assim, o termo aponta para longe de mínimos agudos (onde as mudanças de gradiente são grandes) e permanece pequeno em regiões planas.
4. Ortogonalização Muon: O gradiente aumentado é alimentado nas etapas padrão de momento e ortogonalização do Muon:
   $$M_k = \mu M_{k-1} + \tilde{G}_k$$
   $$O_k = \text{Newton-Schulz}(M_k)$$
   $$W_{k+1} = W_k - \eta (\gamma O_k + \lambda W_k)$$

Fundamentação do Design

Aplicar o termo de aceleração antes da ortogonalização é crítico. Se aplicado após, interromperia a estrutura ortogonal da qual o Muon depende para regularização da norma espectral. Ao transformar o gradiente bruto em uma direção consciente da curvatura primeiro, o buffer de momento captura essas direções enriquecidas, que são então processadas através do pipeline padrão do Muon.

Para parâmetros de valor vetorial (por exemplo, embeddings, vieses), o MONA usa o AdamW como padrão, aderindo às convenções existentes do Muon.

3. Análise Teórica

O artigo fornece uma análise de convergência para o MONA sob suposições padrão (suavidade L, gradientes não viciados com variância limitada e alinhamento direcional esperado).

• Taxa de Convergência: Os autores provam que o MONA atinge uma taxa de convergência de $O(K^{-1/2})$ em cenários não convexos, igualando as taxas assintóticas do SGD e do AdamW. O termo de aceleração afeta as constantes na limitação, mas não a taxa assintótica.
• Escape de Mínimo Agudo: A análise formaliza como o termo de aceleração resiste seletivamente a mínimos agudos. Perto de um ponto estacionário, a aceleração esperada $E[A_k]$ é proporcional à soma ponderada pelo Hessiano das direções passadas. Para direções com autovalores grandes (mínimos agudos), a magnitude da aceleração é grande, empurrando o otimizador para longe. Para autovalores pequenos (direções planas), a aceleração é pequena, permitindo convergência estável. Esse mecanismo preserva a regularização da norma espectral do Muon enquanto adiciona consciência da curvatura.

4. Resultados Empíricos

O MONA foi avaliado em três escalas de modelos de pré-treinamento MoE (1B, 6B e 68B parâmetros) e em uma etapa de ajuste fino supervisionado (SFT).

Desempenho de Pré-treinamento

• Modelos: MOE-1B-A0d2B (400B tokens), MOE-6B-A0d5B (1,2T tokens) e MOE-68B-A3B (700B tokens).
• Métricas: A perda de validação foi monitorada nos domínios de código, raciocínio matemático, inglês geral e texto acadêmico chinês.
• Descobertas: O MONA consistentemente alcançou perda de validação mais baixa que tanto o Muon quanto o AdamW em todas as três escalas de modelo e em todos os domínios. A diferença de desempenho foi particularmente evidente no modelo maior de 68B, onde o MONA demonstrou uma convergência mais estável.

Benchmarks de Tarefas Subsequentes (MOE-68B-A3B)

O modelo de 68B treinado com MONA foi avaliado em uma suíte abrangente de benchmarks:

• Capacidade Geral: O MONA alcançou a pontuação média mais alta (0,4557) em benchmarks incluindo MMLU, MMLU-Pro, CMMLU, CEVAL e BBH, superando o Muon (0,4478) e o AdamW (0,4382).
• Raciocínio Matemático e Geração de Código: O MONA mostrou desempenho superior em GSM8K, GPQA e vários benchmarks de geração de código (HumanEval+, MBPP+, BigCodeBench, LiveCodeBench). Notavelmente, o MONA entregou ganhos significativos em tarefas que exigem representações de código transferíveis.
• Ajuste Fino Supervisionado (SFT): Quando o checkpoint pré-treinado com MONA passou por SFT em dados de código, manteve uma perda de treinamento mais baixa durante todo o processo em comparação com a linha de base pré-treinada com Muon. Avaliações pós-SFT na suíte BigCode confirmaram que modelos inicializados com MONA alcançaram pontuações mais altas em 6 de 7 tarefas em comparação com modelos inicializados com Muon.

Eficiência e Implantação

Para abordar a sobrecarga de memória de armazenar o gradiente anterior ($G_{k-1}$) e o buffer de aceleração ($A_k$), os autores introduziram o MONA-Lite:

• Quantização: Armazenar estados auxiliares em bfloat16 (BF16) em vez de float32.
• Computação em Fluxo: Calcular diferenças de gradiente no local para eliminar a necessidade de um buffer dedicado $G_{k-1}$.
• Resultado: Essas estratégias reduzem a sobrecarga de memória extra em aproximadamente 75% sem sacrificar a qualidade do treinamento. A análise de sobrecarga computacional mostrou que o MONA é apenas ~1% mais lento no nível do passo do otimizador, sem diferença prática no tempo de iteração de ponta a ponta em comparação com o Muon.

5. Significado e Alegações

O artigo alega que o MONA conecta com sucesso a lacuna entre a eficiência geométrica do Muon e a necessidade de aceleração consciente da curvatura. Seu significado principal reside em:

1. Integração Algorítmica: Demonstrar que termos de aceleração podem ser integrados em estruturas de ortogonalização de matrizes sem interromper suas propriedades geométricas.
2. Escalabilidade: Validar o método em uma ampla gama de tamanhos de modelo (1B a 68B parâmetros) e contagens de tokens (até 1 trilhão de tokens), mostrando melhorias consistentes sobre tanto o AdamW quanto o Muon.
3. Generalização: Fornecer evidências empíricas de que escapar de mínimos agudos durante o pré-treinamento leva a uma melhor generalização em tarefas subsequentes, particularmente em raciocínio complexo e geração de código.
4. Praticidade: Oferecer uma variante eficiente em memória (MONA-Lite) que torna o otimizador viável para ambientes com recursos limitados.

Os autores concluem que o MONA representa um otimizador de última geração para treinamento de LLMs em grande escala, alcançando convergência superior e desempenho em tarefas subsequentes, mantendo a eficiência computacional da estrutura do Muon.