Delving into Muon and Beyond: Deep Analysis and Extensions

Resumo Técnico: Aprofundando no Muon e Além: Análise Detalhada e Extensões

Declaração do Problema
O otimizador Muon ganhou atenção recentemente por seu desempenho empírico no treinamento de grandes modelos de linguagem, utilizando atualizações ortogonalizadas em nível de matriz. No entanto, os mecanismos subjacentes do Muon e sua relação com otimizadores adaptativos estabelecidos, como o Adam, permanecem insuficientemente compreendidos. A literatura existente frequentemente confunde os efeitos do Muon com técnicas auxiliares (ex: QK-Norm, QK-Clip) ou carece de uma comparação sistemática contra baselines padrão. Além disso, não está claro se a ortogonalização espectral do Muon representa uma estratégia de otimização universalmente superior ou uma forma específica de estabilização, e se ele pode ser melhorado ou generalizado.

Metodologia
Os autores propõem um framework espectral unificado para analisar o Muon e suas variantes. Eles definem o Muon como o endpoint $p=0$ de uma família de transformações espectrais da forma $U\Sigma^p V^\top$, onde $U\Sigma V^\top$ é a Decomposição de Valores Singulares (SVD) da matriz de atualização. O estudo investiga expoentes intermediários $p \in \{1, 1/2, 1/4, 0\}$:

• $p=1$: Identidade (atualizações padrão).
• $p=1/2, 1/4$: Compressão espectral parcial.
• $p=0$: Ortogonalização total (Muon).

Essas transformações são aplicadas a dois tipos distintos de entrada:

1. Momento do primeiro momento ($O_{mom}$): Equivalente aos inputs de momentum SGD.
2. Atualizações normalizadas por RMS ($O_{rms}$): Equivalente aos inputs do Adam, onde os gradientes são normalizados pelo segundo momento.

Para permitir o cálculo eficiente de potências fracionárias ($p=1/2, 1/4$) sem SVD explícita, os autores desenvolveram uma iteração de Newton-Schulz acoplada. Este algoritmo calcula simultaneamente a raiz quadrada da matriz e sua inversa usando apenas multiplicações de matrizes, garantindo estabilidade numérica e escalabilidade.

A avaliação experimental é conduzida em um modelo GPT-2 de 124M de parâmetros treinado no OpenWebText. O estudo utiliza um setup altamente controlado:

• Taxas de Aprendizado Desacopladas: Taxas de aprendizado separadas são usadas para parâmetros de matriz e vetores para isolar o efeito do otimizador.
• Sem Técnicas Auxiliares: Weight decay, QK-Norm e QK-Clip são desativados para garantir comparações justas.
• Ajuste Sistemático: Uma busca de taxa de aprendizado em duas etapas (varredura logarítmica grosseira seguida de refinamento fino) é realizada para todas as oito variantes de otimizadores (4 expoentes espectrais $\times$ 2 tipos de input).

Principidades Contribuições

1. Framework Espectral Unificado: O artigo introduz um framework formal $U\Sigma^p V^\top$ que posiciona o Muon como um caso específico ($p=0$) e propõe variantes intermediárias ($p=1/2, 1/4$) para interpolar entre atualizações padrão e ortogonalização total.
2. Algoritmo Eficiente: Uma iteração de Newton-Schulz acoplada é desenvolvida para calcular atualizações espectrais fracionárias ($p=1/2, 1/4$) de forma eficiente, evitando o custo computacional da SVD explícita.
3. Comparação Empírica Rigorosa: O estudo fornece uma comparação controlada entre Muon, Adam e suas variantes espectrais, isolando efeitos espectrais de fatores de confusão como weight decay ou técnicas de normalização auxiliar.

Principais Resultados
Os experimentos produzem três descobertas primárias:

1. Estabilização de Atualizações de Primeiro Momento: Em inputs de primeiro momento (mSGD), o Muon ($p=0$) fornece uma estabilização significativa, permitindo taxas de aprendizado muito mais altas e maior robustez em comparação ao mSGD padrão. No entanto, a compressão parcial ($p=1/4$) frequentemente alcança um desempenho final melhor uma vez que a estabilidade é estabelecida.
2. Ganhos Limitados em Inputs Normalizados por RMS: Quando aplicado a inputs normalizados por RMS (estilo Adam), a compressão espectral produz retornos decrescentes. A ortogonalização total (AdamZ/Muon aplicado a $O_{rms}$) não supera o Adam padrão e pode ser inferior à compressão parcial. Isso sugere que a normalização RMS elemento a elemento já mitiga grande parte da anisotropia que a compressão espectral visa combater.
3. Superioridade de Métodos de Segundo Momento: Em todos os expoentes espectrais, otimizadores que utilizam normalização de segundo momento (baseada em RMS) superam consistentemente seus equivalentes de primeiro momento (apenas momentum) em termos de estabilidade e perda final.

Significância e Alegações
O artigo conclui que o Muon é melhor compreendido como uma forma eficaz de normalização espectral em vez de um método de otimização universalmente superior.

• Estabilização vs. Desempenho: A principal contribuição do Muon é estabilizar as atualizações de primeiro momento ao descartar as magnitudes dos valores singulares, o que evita que direções com grandes valores singulares dominem. No entanto, este mecanismo de "achatamento" é bruto; ele remove o filtro baseado em magnitude, potencialmente amplificando direções de gradiente ruidosas (pequenos valores singulares) e subestimando modos dominantes informativos.
• Utilidade Contextual: Embora a compressão espectral ofereça benefícios reais de estabilização, especialmente quando o input de atualização não é já normalizado, a ortogonalização total ($p=0$) não é um substituto para os modernos otimizadores de segundo momento normalizados por RMS, como o Adam. Os autores argumentam que as vantagens relatadas do Muon em trabalhos anteriores podem ser atribuídas à combinação específica de técnicas utilizadas (ex: QK-Norm) em vez da ortogonalização em si.

O trabalho enfatiza que, embora as transformações espectrais sejam ferramentas poderosas para controlar a anisotropia, elas não superam inerentemente os benefícios da escala adaptativa de segundo momento em regimes de treinamento de larga escala.