How Much Orthogonalization Does Muon Need?

Resumo Técnico: Quanta Ortogonalização o Muon Precisa?

Definição do Problema
O otimizador Muon melhora o treinamento de redes neurais ao substituir atualizações de momentum mal condicionadas por atualizações aproximadamente semi-ortogonais. Enquanto a álgebra linear numérica padrão visa uma decomposição polar de alta precisão (mapeando todos os valores singulares para 1), o contexto do Muon sugere que a precisão exata pode ser desnecessária. A questão central abordada por este trabalho é: Quanta precisão de ortogonalização o Muon realmente exige para preservar a qualidade do treinamento? Especificamente, o artigo investiga se o custo computacional de iterações polinomiais de alto grau (usadas em métodos existentes como Muon-Jordan e Polar Express) é justificado, ou se uma aproximação mais barata e de menor grau pode alcançar resultados de treinamento comparáveis.

Metodologia
Os autores propõem um esquema de Newton–Schulz cúbico relaxado (denominado cubic5) projetado especificamente para a banda de valores singulares de baixa precisão (bfloat16) do Muon.

1. Alvos Relaxados: Ao contrário da decomposição polar clássica que visa a precisão de máquina, este método visa uma banda relaxada de $[0,7, 1,3]$ para os valores singulares. Ele reconhece que, em bfloat16, valores singulares abaixo de um limite inferior efetivo ($l_0 = 7 \times 10^{-3}$) não são praticamente significativos.
2. Derivação de Coeficientes Adaptativos: Os autores derivam um polinômio cúbico adaptativo $f_t(x) = a_t x + b_t x^3$ em cada iteração. Os coeficientes são calculados com base no atual limite inferior de pior caso ($l_t$) e um pico superior fixo ($u=1,3$). Isso garante que o polinômio mapeie o intervalo protegido $[l_t, r_t]$ para a banda alvo enquanto equilibra as extremidades para maximizar o novo limite inferior.
3. Eficiência Computacional:
   • Cubic5: Um único passo cúbico requer duas multiplicações de matrizes dominantes (formando a matriz de Gram e multiplicando-a de volta). Cinco passos exigem 10 multiplicações de matrizes dominantes.
   • Baseline (Muon-Jordan/Polar Express): Um passo quíntico requer três multiplicações de matrizes dominantes. Cinco passos exigem 15 multiplicações de matrizes dominantes.
   • O método utiliza o "truque da transposta" para matrizes altas para minimizar as operações na matriz de Gram menor.

Principais Contribuições

• Derivação de um Esquema Cúbico Relaxado: O artigo deriva um esquema Newton–Schulz cúbico de cinco passos com coeficientes adaptados ao limite inferior efetivo do bfloat16, alcançando a banda espectral alvo em cinco iterações.
• Caracterização da Fronteira Custo-Qualidade: O trabalho estabelece que a qualidade do treinamento não é governada monotonicamente pela precisão da decomposição polar. Ele demonstra que uma aproximação cúbica de menor custo pode alcançar paridade de desempenho com métodos quínticos de maior custo e até mesmo com o SVD exato em regimes específicos.
• Validação Empírica em Diferentes Escalas: Os autores validam o método através de:
  • Diagnósticos Sintéticos: Confirmando que a transferência do polinômio escalar se mantém na aritmética de matriz de baixa precisão.
  • Ablações NanoGPT: Comparando cubic5, Polar Express truncado, Muon-Jordan truncado e SVD explícito em FP32 no GPT-2 Small.
  • Treinamento em Larga Escala: Treinamento de modelos híbridos MoE/Mamba (Nemotron-3-Nano) variando de 1B a 4B de parâmetros usando Megatron-LM.

Resultos

• Benchmarks Sintéticos: O método cubic5 mapeia com sucesso valores singulares no intervalo protegido $[7 \times 10^{-3}, 1]$ para a banda alvo sem explosão numérica. Microbenchmarks em GPUs NVIDIA H200 mostram que o cubic5 reduz o custo da sub-rotina de ortogonalização em relação ao Polar Express, com acelerações variando de 1,16x a 1,52x dependendo da forma da matriz (retangular vs. quadrada).
• Ablações GPT-2 Small:
  • O cubic5 de cinco passos atinge uma perda de validação final de 3,2811.
  • O Polar Express de cinco passos atinge 3,2798.
  • O Muon-Jordan de cinco passos atinge 3,2822.
  • O SVD explícito em FP32 atinge 3,2809.
  • Observação: A lacuna entre as variantes mais fortes é inferior a 0,003 na perda de validação. Notavelmente, o SVD exato não supera os métodos polinomiais aproximados, e adicionar um quinto passo ao Muon-Jordan degradou ligeiramente o desempenho em comparação com quatro passos.
• Escalabilidade Híbrida MoE/Mamba (1B–4B):
  • Nos modelos de 1B e 2B, o cubic5 superou ligeiramente o Muon-Jordan (perda de validação menor por $\approx 10^{-3}$).
  • No modelo de 4B, o cubic5 foi ligeiramente superior ($\approx 10^{-3}$), mas as trajetórias de treinamento foram quase indistinguíveis após a fase inicial.
  • A diferença relativa na perda de treinamento diminuiu rapidamente, flutuando em torno de zero pelo restante do orçamento de treinamento.

Significância e Alegações
O artigo alega modestamente que o cubic5 é uma variante prática e de baixo custo para a ortogonalização do Muon, em vez de uma regra de atualização universalmente superior.

• Otimizador vs. Precisão de Decomposição: Os resultados apoiam a hipótese de que o Muon se beneficia de um efeito de "remodelagem espectral" em vez de uma precisão estrita de decomposição polar. O mapeamento exato dos valores singulares que otimiza o treinamento depende da fase de treinamento e da taxa de aprendizado, não apenas da convergência matemática para 1.
• Utilidade Prática: O cubic5 oferece um ponto principiado na fronteira custo-qualidade. Ele reduz a contagem de multiplicações de matrizes dominantes em cerca de um terço (10 vs. 15) enquanto mantém a paridade da qualidade de treinamento empírica com a atualização quíntica padrão do Muon-Jordan nos regimes testados.
• Limitações: Os autores declaram explicitamente que estes resultados estabelecem a paridade de qualidade de treinamento, não acelerações de throughput de ponta a ponta. A aceleração real de tempo de execução (wall-clock) depende de kernels de hardware, batching e da fração de tempo que a etapa de ortogonalização ocupa dentro do loop completo do otimizador. O artigo não afirma que o cubic5 é um algoritmo de decomposição polar melhor que o Polar Express do ponto de vista matemático, mas sim uma transformação espectral mais eficiente para o contexto específico do treinamento do Muon.