MARS-M: Cuando la reducción de varianza se encuentra con las matrices

MARS-M: When Variance Reduction Meets Matrices

October 20, 2025

Autores: Yifeng Liu, Angela Yuan, Quanquan Gu

cs.AI

Resumen

Los optimizadores precondicionados basados en matrices, como Muon, han demostrado recientemente ser más eficientes que los optimizadores basados en escalares para entrenar redes neuronales a gran escala, incluidos los grandes modelos de lenguaje (LLMs). Por otro lado, evaluaciones recientes de optimizadores para el preentrenamiento de LLMs han demostrado que técnicas de reducción de varianza como MARS pueden lograr aceleraciones sustanciales respecto a optimizadores estándar que no emplean reducción de varianza. En este artículo, para lograr lo mejor de ambos mundos, presentamos MARS-M, un nuevo optimizador que integra la técnica de reducción de varianza de MARS con Muon. Bajo condiciones de regularidad estándar, demostramos que Muon-M converge a un punto estacionario de primer orden con una tasa de \(\mathcal{O}(T^{-1/3})\), lo que mejora la tasa de \(\mathcal{O}(T^{-1/4})\) alcanzada por Muon. Nuestros resultados empíricos en tareas de modelado de lenguaje y visión por computadora demuestran que MARS-M produce consistentemente pérdidas más bajas y un rendimiento mejorado en varios puntos de referencia posteriores. La implementación de MARS-M está disponible en 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.