MARS-M: Quando la riduzione della varianza incontra le matrici

MARS-M: When Variance Reduction Meets Matrices

October 20, 2025

Autori: Yifeng Liu, Angela Yuan, Quanquan Gu

cs.AI

Abstract

Gli ottimizzatori precondizionati basati su matrici, come Muon, hanno recentemente dimostrato di essere più efficienti degli ottimizzatori scalari per l'addestramento di reti neurali su larga scala, inclusi i grandi modelli linguistici (LLM). D'altro canto, recenti benchmark sugli ottimizzatori per il pre-addestramento di LLM hanno evidenziato che tecniche di riduzione della varianza come MARS possono ottenere accelerazioni sostanziali rispetto agli ottimizzatori standard che non impiegano tale riduzione. In questo articolo, per unire i vantaggi di entrambi gli approcci, introduciamo MARS-M, un nuovo ottimizzatore che integra la tecnica di riduzione della varianza di MARS con Muon. Sotto le consuete condizioni di regolarità, dimostriamo che Muon-M converge verso un punto stazionario del primo ordine con un tasso di \(\mathcal{O}(T^{-1/3})\), migliorando il tasso \(\mathcal{O}(T^{-1/4})\) raggiunto da Muon. I nostri risultati empirici su attività di modellazione linguistica e visione artificiale mostrano che MARS-M produce costantemente loss inferiori e prestazioni migliorate su vari benchmark downstream. L'implementazione di MARS-M è disponibile all'indirizzo 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.