Adam Aprimora Muon: Estimação de Momento Adaptativo com Momento Ortogonalizado

Adam Improves Muon: Adaptive Moment Estimation with Orthogonalized Momentum

February 19, 2026

Autores: Minxin Zhang, Yuxuan Liu, Hayden Scheaffer

cs.AI

Resumo

A otimização estocástica eficiente normalmente integra uma direção de atualização que apresenta bom desempenho no regime determinístico com um mecanismo de adaptação a perturbações estocásticas. Enquanto o Adam utiliza estimativas adaptativas de momento para promover estabilidade, o Muon aproveita a estrutura matricial das camadas de pesos através do momento ortogonalizado, demonstrando desempenho superior no treinamento de modelos de linguagem grandes. Propomos um novo otimizador e uma extensão diagonal, NAMO e NAMO-D, fornecendo a primeira integração fundamentada do momento ortogonalizado com a adaptação ao ruído do tipo Adam baseada em norma. O NAMO dimensiona o momento ortogonalizado usando um único tamanho de passo adaptativo, preservando a ortogonalidade enquanto supera o Muon com custo adicional insignificante. O NAMO-D, por sua vez, multiplica à direita o momento ortogonalizado por uma matriz diagonal com entradas limitadas. Este projeto permite a adaptação ao ruído por neurônio e alinha-se com a estrutura hessiana comum quase bloco-diagonal. Sob premissas padrão, estabelecemos taxas de convergência ótimas para ambos os algoritmos no cenário determinístico e mostramos que, no cenário estocástico, suas garantias de convergência se adaptam ao nível de ruído dos gradientes estocásticos. Experimentos com pré-treinamento de modelos GPT-2 demonstram um desempenho melhorado tanto do NAMO quanto do NAMO-D em comparação com as baselines AdamW e Muon, com o NAMO-D alcançando ganhos adicionais sobre o NAMO através de um hiperparâmetro de limitação adicional que equilibra os objetivos concorrentes de manter uma direção de atualização bem condicionada e alavancar a adaptação ao ruído em granularidade fina.

English

Efficient stochastic optimization typically integrates an update direction that performs well in the deterministic regime with a mechanism adapting to stochastic perturbations. While Adam uses adaptive moment estimates to promote stability, Muon utilizes the weight layers' matrix structure via orthogonalized momentum, showing superior performance in large language model training. We propose a new optimizer and a diagonal extension, NAMO and NAMO-D, providing the first principled integration of orthogonalized momentum with norm-based Adam-type noise adaptation. NAMO scales orthogonalized momentum using a single adaptive stepsize, preserving orthogonality while improving upon Muon at negligible additional cost. NAMO-D instead right-multiplies orthogonalized momentum by a diagonal matrix with clamped entries. This design enables neuron-wise noise adaptation and aligns with the common near block-diagonal Hessian structure. Under standard assumptions, we establish optimal convergence rates for both algorithms in the deterministic setting and show that, in the stochastic setting, their convergence guarantees adapt to the noise level of stochastic gradients. Experiments on pretraining GPT-2 models demonstrate improved performance of both NAMO and NAMO-D compared to the AdamW and Muon baselines, with NAMO-D achieving further gains over NAMO via an additional clamping hyperparameter that balances the competing goals of maintaining a well-conditioned update direction and leveraging fine-grained noise adaptation.