MARS-M: Quando a Redução de Variância Encontra as Matrizes

MARS-M: When Variance Reduction Meets Matrices

October 20, 2025

Autores: Yifeng Liu, Angela Yuan, Quanquan Gu

cs.AI

Resumo

Os otimizadores pré-condicionados baseados em matrizes, como o Muon, demonstraram recentemente maior eficiência do que os otimizadores baseados em escalares para o treinamento de redes neurais em larga escala, incluindo grandes modelos de linguagem (LLMs). Por outro lado, benchmarks recentes sobre otimizadores para pré-treinamento de LLMs demonstraram que técnicas de redução de variância, como o MARS, podem alcançar acelerações substanciais em relação a otimizadores padrão que não empregam redução de variância. Neste artigo, para obter o melhor dos dois mundos, introduzimos o MARS-M, um novo otimizador que integra a técnica de redução de variância do MARS com o Muon. Sob condições de regularidade padrão, provamos que o Muon-M converge para um ponto estacionário de primeira ordem a uma taxa de \(\mathcal{O}(T^{-1/3})\), o que representa uma melhoria em relação à taxa de \(\mathcal{O}(T^{-1/4})\) atingida pelo Muon. Nossos resultados empíricos em tarefas de modelagem de linguagem e visão computacional demonstram que o MARS-M produz consistentemente perdas menores e desempenho melhorado em vários benchmarks *downstream*. A implementação do MARS-M está disponível em https://github.com/AGI-Arena/MARS/MARS_M.

English

Matrix-based preconditioned optimizers, such as Muon, have recently been shown to be more efficient than scalar-based optimizers for training large-scale neural networks, including large language models (LLMs). On the other hand, recent benchmarks on optimizers for LLM pre-training have demonstrated that variance-reduction techniques such as MARS can achieve substantial speedups over standard optimizers that do not employ variance reduction. In this paper, to achieve the best of both worlds, we introduce MARS-M, a new optimizer that integrates the variance reduction technique in MARS with Muon. Under standard regularity conditions, we prove that Muon-M converges to a first-order stationary point at a rate of mathcal{O}(T^{-1/3}), which improves upon mathcal{O}(T^{-1/4}) rate attained by Muon. Our empirical results on language modeling and computer vision tasks demonstrate that MARS-M consistently yields lower losses and improved performance across various downstream benchmarks. The implementation of MARS-M is available at https://github.com/AGI-Arena/MARS/MARS_M.