Resumo Técnico: Muown - Controle de Norma de Linha para Otimização Muon

Declaração do Problema

O Muon emergiu como um otimizador competitivo para o pré-treinamento de modelos de linguagem de grande escala (LLM), frequentemente superando o AdamW ao explorar a estrutura matricial das camadas lineares do Transformer. No entanto, observações recentes indicam que o comportamento do Muon em escala é sensível ao decaimento de pesos e propenso a instabilidades numéricas, como transbordamentos em bfloat16 e logits de atenção explodindo.

Uma questão crítica identificada é o desvio ascendente consistente da norma espectral ($\|W_t\|_{S_\infty}$) das matrizes de pesos durante o treinamento quando o Muon é utilizado sem decaimento de pesos desacoplado. Esse desvio viola a condição espectral necessária para aprendizado de características estável ($\|W\|_{S_\infty} = \Theta(\sqrt{m/n})$). Embora o decaimento de pesos desacoplado possa mitigar isso ao reduzir uniformemente a matriz de pesos, trata-se de uma intervenção não direcionada que escala todo o espectro independentemente de quais direções específicas impulsionam o crescimento, potencialmente retardando a convergência nas fases iniciais do treinamento. O artigo busca identificar o motor empírico específico desse desvio da norma espectral e propor um mecanismo de controle mais direcionado.

Metodologia: Muown

Os autores introduzem o Muown, um substituto direto para o Muon que aborda o desvio da norma espectral por meio de uma estratégia de reparametrização em vez de restrições agregadas.

1. Decomposição da Norma Espectral

A ideia central é uma decomposição da norma espectral em dois fatores (Proposição 1):
$$\|W_t\|_{S_\infty}^2 = \|g_t\|_\infty^2 \cdot \lambda_{\max}(P_t C_t P_t)$$
Onde:

• $g_t = \|W_t\|_{row}$ é o vetor de magnitudes de linha (normas euclidianas das linhas).
• $D_t$ é a matriz normalizada por linha ($W_t = \text{Diag}(g_t)D_t$).
• $C_t = D_t D_t^\top$ é a matriz de Gram das linhas de norma unitária.
• $P_t$ é uma matriz diagonal de magnitudes de linha normalizadas.
• $\lambda_{\max}(P_t C_t P_t)$ representa o fator de coerência de linha.

A análise empírica em modelos de 500M revela que o fator de magnitude de linha ($\|g_t\|_\infty$) é o principal motor do desvio da norma espectral, movendo-se em sincronia com $\|W_t\|_{S_\infty}$. Em contraste, o fator de coerência de linha permanece bem comportado e limitado durante todo o treinamento.

2. Reparametrização e Otimização

Com base nesse diagnóstico, o Muown reparametriza a matriz de pesos $W$ como:
$$W(g, R) = \text{Diag}\left(\frac{g}{\|R\|_{row}}\right) R$$
Aqui, $g \in \mathbb{R}^m$ representa o vetor explícito de magnitudes de linha, e $R \in \mathbb{R}^{m \times n}$ representa o componente de direção.

• Componente de Direção ($R$): O Muon é aplicado sem alterações a $R$, tratando-o sob a geometria da norma espectral (descida espectral mais íngreme).
• Componente de Magnitude ($g$): As magnitudes de linha $g$ são promovidas a variáveis explícitas do otimizador. Elas são atualizadas sob a geometria $\ell_\infty$, que é a norma dual natural para a norma do operador de matrizes diagonais.
• Regra de Atualização: A atualização para $g$ utiliza um otimizador Adam (atuando como um SignSGD suavizado sob a geometria $\ell_\infty$) para lidar com ruído estocástico, enquanto a atualização para $R$ segue o procedimento padrão do Muon com ortogonalização Newton-Schulz.

Essa abordagem é implementada implicitamente dentro do estado do otimizador, deixando a passagem direta da arquitetura do modelo inalterada.

Contribuições Principais

1. Diagnóstico do Desvio Espectral: O artigo decompõe a norma espectral para identificar a magnitude máxima de linha ($\|g_t\|_\infty$) como o motor empírico do crescimento da norma espectral sob o Muon, distinguindo-a do fator de coerência de linha estável.
2. Algoritmo Muown: Um novo otimizador que trata as magnitudes de linha como variáveis treináveis explícitas atualizadas sob a geometria $\ell_\infty$, preservando as atualizações direcionais do Muon. Isso evita a redução indiscriminada do decaimento de pesos e os priores de ortogonalidade das reparametrizações espectrais.
3. Garantias de Convergência: Os autores provam que o Muown atinge taxas de convergência não convexas ótimas em regimes determinísticos e estocásticos. A análise é conduzida em uma norma dual mista ($\ell_1$ para magnitudes, Schatten-1 para direções) alinhada com as geometrias de atualização. O limite teórico sugere uma redução no coeficiente de ruído estocástico efetivo em comparação com o Muon.
4. Validação Empírica: Experimentos extensivos no FineWeb-Edu, em tamanhos de modelo de 124M a 2,7B de parâmetros, demonstram melhorias consistentes.

Resultados Experimentais

• Melhorias na Perplexidade: O Muown alcança consistentemente perplexidade mais baixa do que Muon, SOAP, AdamW e Lion em todas as escalas de modelo testadas (124M a 2,7B) e taxas de aprendizado.
• Robustez a Hiperparâmetros:
  • Taxa de Aprendizado: O Muown alarga o platô de taxas de aprendizado quase ótimas em diferentes larguras de modelo, reduzindo a sensibilidade ao ajuste de taxa de aprendizado por escala.
  • Decaimento de Pesos: O Muown alivia a necessidade de ajuste de decaimento de pesos; ele performa bem sem decaimento de pesos ($\lambda=0$), enquanto o Muon requer ajuste cuidadoso (ótimo $\lambda=0,1$) para evitar divergência ou desempenho subótimo.
• Estabilidade: O Muown previne o desvio sistemático da norma espectral observado no Muon. Em regimes de alta taxa de aprendizado (ex: $4 \times 10^{-3}$), o Muon frequentemente diverge, enquanto o Muown permanece estável e próximo do ótimo.
• Propriedades do Gradiente: A análise mostra que o Muown mantém um rank efetivo normalizado mais alto dos gradientes de minibatch e exibe um coeficiente de ruído estocástico mais baixo durante todo o treinamento em comparação com o Muon.
• Sobrecarga Computacional: O método introduz sobrecarga de tempo de etapa negligenciável. Quando o otimizador é particionado entre ranks (ambiente distribuído), os tempos de etapa para Muown e Muon tornam-se indistinguíveis. A sobrecarga de memória é mínima (linear no número de linhas), adicionando apenas alguns megabytes por camada.

Significado e Alegações

O artigo alega que o desvio da norma espectral observado no Muon não é uma propriedade irredutível do otimizador, mas uma consequência do crescimento descontrolado da magnitude de linha. Ao isolar esse motor e promovê-lo a uma variável explícita do otimizador, o Muown alcança:

• Convergência Melhorada: Melhor perplexidade e estabilidade sem a necessidade de decaimento de pesos agressivo.
• Fundamentação Teórica: Taxas de convergência ótimas em um framework consciente da geometria que leva explicitamente em conta a natureza mista das atualizações.
• Eficiência Prática: Um substituto direto que iguala o tempo de execução do Muon em ambientes distribuídos, oferecendo ao mesmo tempo robustez superior na escolha de hiperparâmetros.

Os autores posicionam o Muown como uma solução que preenche a lacuna entre os requisitos teóricos de estabilidade do aprendizado de características e as necessidades práticas do treinamento de LLMs em grande escala, oferecendo uma alternativa mais direcionada ao decaimento de pesos desacoplado.