MARS: Unleashing the Power of Variance Reduction for Training Large Models
Authors: Huizhuo Yuan
Yifeng Liu
Shuang Wu
zhou Xun
Quanquan Gu
Venue: ICML 2025 poster (2025)
Keywords: variance reduction
adaptive learning
large models
Primary area: optimization

Abstract: Training deep neural networks--and more recently, large models--demands efficient and scalable optimizers. Adaptive gradient algorithms like Adam, AdamW, and their variants have been central to this task. Despite the development of numerous variance reduction algorithms in the past decade aimed at accelerating stochastic optimization in both convex and nonconvex settings, variance reduction has not found widespread success in training deep neural networks or large language models. Consequently, it has remained a less favored approach in modern AI. In this paper, to unleash the power of variance reduction for efficient training of large models, we propose a unified optimization framework, MARS (**M**ake v**A**riance **R**eduction **S**hine), which reconciles preconditioned gradient methods with variance reduction via a scaled stochastic recursive momentum technique. Within our framework, we introduce three instances of MARS that leverage preconditioned gradient updates based on AdamW, Lion, and Shampoo, respectively. We also draw a connection between our algorithms and existing optimizers. Experimental results on training GPT-2 models indicate that MARS consistently outperforms AdamW by a large margin.

MARS: Unleashing the Power of Variance Reduction for Training Large Models
Huizhuo Yuan + 12 Yifeng Liu
Shuang Wu?
Xun Zhou? Quanquan Gu
Abstract
trained with adaptive gradient methods. Numerous efforts have been made to improve adaptive gradient methods from both first-order and second-order optimization perspectives. For example, You et al. (2019) introduced LAMB, a lay- erwise adaptation technique that boosts training efficiency for BERT (Devlin, 2018). Using symbolic search, Chen et al. (2023) developed Lion, achieving faster training and reduced memory usage. Liu et al. (2023) designed Sophia, leveraging stochastic diagonal Hessian estimators to acceler- ate training. Gupta et al. (2018) proposed Shampoo, which performs stochastic optimization over tensor spaces with preconditioning matrices for each dimension. Anil et al. (2020) further refined Shampoo to give a scalable, practical version. Recently, Vyas et al. (2024) showed that Sham- poo is equivalent to Adafactor (Shazeer & Stern, 2018) in the eigenbasis of Shampoo's 's preconditioner, and introduced SOAP, which stabilizes Shampoo with Adam. Importantly, recent studies (Kaddour et al., 2024; Zhao et al., 2024) have shown that these optimizers perform on par with AdamW in LLM pretraining, yet do not outperform it. This suggests the ongoing challenge developing adaptive gradient meth- ods superior to Adam and AdamW for large-scale model training.
Training deep neural networks-and more re- cently, large models demands efficient and scal- able optimizers. Adaptive gradient algorithms like Adam, AdamW, and their variants have been central to this task. Despite the development of numerous variance reduction algorithms in the past decade aimed at accelerating stochastic opti- mization in both convex and nonconvex settings, variance reduction has not found widespread suc- cess in training deep neural networks or large language models. Consequently, it has remained less favored approach in modern AI. In this pa- per, to unleash the power of variance reduction for efficient training of large models, we propose unified optimization framework, MARS (Make vAriance Reduction Shine), which reconciles pre- conditioned gradient methods with variance reduc- tion via a scaled stochastic recursive momentum technique. Within our framework, we introduce three instances of MARS that leverage precondi- tioned gradient updates based on AdamW, Lion, and Shampoo, respectively. We also draw con- nection between our algorithms and existing opti- mizers. Experimental results on training GPT-2 models indicate that MARS consistently outper- forms Adam W by a large margin.
Since adaptive gradient methods face challenges of high stochastic gradient variance, and language model training inherently involves high-variance optimization problem (McCandlish al., 2018), it is natural to consider variance reduction techniques to address this challenge. There exists a large body of literature on variance reduction for stochas- tic optimization, such as SAG (Roux et al., 2012), SVRG (Johnson & Zhang, 2013), STORM (Cutkosky & Orabona, 2019), which can improve the convergence of stochastic optimization. However, variance reduction has not found widespread success in training deep neural networks large language models. Defazio & Bottou (2019) discussed why variance reduction can be ineffective in deep learning due to factors such as data augmentation, batch normalization, and dropout, which disrupt the finite-sum structure required by variance reduction principles. Nevertheless, in training language models, data augmentation, batch normalization, and dropout nowadays rarely used, which opens the door for applying variance reduction techniques in optimizing these models. This naturally leads to the following research question:
Introduction
Adaptive gradient methods such as Adam (Kingma & Ba, 2015) and AdamW (Loshchilov & Hutter, 2019) have be- come the predominant optimization algorithms in deep learning. With the surge of large language models, the majority of the renowned models, including GPT-2 (Rad- ford et al., 2019), GPT-3 (Brown, 2020), PaLM (Chowd- hery et al., 2023) and Llama (Dubey et al., 2024) are
'Equal contribution Department of Computer Science, Univer- sity of California, Los Angeles, California, USA (This work was done during Yifeng's internship at ByteDance Seed) BBytDDacce Seed, San Jose, Califomia, USA Dance Seed, Boijing, China. Correspondence to: Quanquan Gu <qgu @aw.@eccc...cuu....
Proceedings of the 4nnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

MARS: Unleashing the Power of Variance Reduction for Training Large Models
Can variance reduction technique be applied to improve the performance of training large models?
sequence of independent random variables which denote the training data for each step. For some objective function f Ex] = F(x) for
In this paper, answer the above question affirmatively by introducing a novel optimization framework called MARS (Make v vAriance Reduction Shine), which incorporates vari- ance reduction into adaptive gradient methods. Notably, we introduce a scaling parameter into the stochastic recur- sive momentum (STORM) (Cutkosky & Orabona, 2019) to adjust the strength of variance reduction and define a new gradient estimator. This gradient estimator undergoes gradient clipping and is subsequently subjected to expo- nential averaging. When the variance reduction strength is set to 1, it recovers the vanilla STORM momentum. In addition, the second-order momentum update is defined by the reweighted intermediate variable. These together ensure optimization stability throughout the training process. We summarize our major contributions of this paper as follows:
that is differentiable, we assume [x] vx, Vt. In our algorithm, the training data of the current step & and previous step Ê are used for attaining different gradient for the same parameter Xx, so we just explicitly indicate these variables for function f.
2 Preliminaries
In this section, we review the preliminaries of stochastic op- timization, including standard stochastic gradient methods and variance reduction.
We consider minimizing an objective function F(+) : IR + R as follows:
min F(x) = [Ex
(2.1)
where f(x. E) is possibly nonconvex loss function, x E Rd is the optimization variable, & is random vector (e.g.. training data point) drawn from an unknown data dis- tribution D. We assume the access to the first-order ora- cle, which returns an unbiased estimator of the gradient f(V $ = VF(x). The standard stochastic gradient descent (SGD) algorithm yields:
We propose  unified framework for preconditioned variance reduction, namely MARS. At its core, MARS comprises two major components: (1) scaled stochastic recursive momentum, which provides a variance-reduced estimator of the full gradient for bet- ter gradient complexity; and (2) the preconditioned up- date, which approximates the second-order Newton's method for better per- itrireraaiinn complexity. By com- bining preconditioned gradient methods with variance reduction, MARS achieves the best of both worlds, ac- celerating the search for critical points in optimization.
X = x, - 7/ V Et),
(2.2)
where 7)t > 0 is the learning rate or step size. SGD needs ((E stochastic gradient evaluations (i.c., gradi- ent complexity or incremental first-order oracle complexity) to find e-approximate first-order stationary points, i.e.,
The MARS framework is versatile, accommodating all existing full matrix or diagonal Hessian approxi- mations. Under this framework, we utilize three dis- tinct designs of the preconditioning matrix, resulting in three specific instances of our MARS framework: MARS-AdamW. MARS-Lion, and MARS MARS-Shampoo. Each variant demonstrates compatibility with their corresponding preconditioning in AdamW, Lion, and Shampoo, showing that MARS can seamlessly inte- grate with and do variance reduction on these estab- lished methods.
F (Ghadimi & Lan, 2013).
To accelerate the convergence of SGD, variance reduction techniques have been extensively researched in both the machine learning and optimization communities over the past decade, resulting in numerous algorithms for convex optimization- as SAG (Roux et al., 2012), SVRG (Johnson & Zhang, 2013). SAGA (Defazio et al., 2014), and SARAH (Nguyen et al., 2017a) as well as for nonconvex optimization, including SVRG (Allen-Zhi & Yuan, 2016; Reddi et al., 2016), SNVRG (Zhou et al., 2020), SPIDER (Fang et al., 2018), and STORM (Cutkosky & Orabona, 2019), among others. Notably, for nonconvex optimization, SNVRG (Zhou et al., 2020). SPIDER (Fang et al., 2018) and STORM (Cutkosky & Orabona, 2019) can improve the gradient complexity of SGD from O(E 4) to O(E ) demonstrating a provable advantage.
Empirically, we evaluated MARS on GPT-2 fine-tuning tasks using the WeenWe dataset. It demonstrates superior performance on GPT-2 large: Adam re- quires 50 billion tokens to reach a validation loss of 2.58, whereas MARS only requires 28 billion tokens, and it achieves final validation loss of 2.51. Further- more, on the downstream task Hellaswag, MARS im- proved accuracy to 44.64%, outperforming AdamW's 41.70% after training 50 billion tokens. And the code is available at https: http:://ggttthcccc com/
At the heart of variance reduction techniques is a variance- reduced stochastic gradient, exemplified by the method pro- posed by Johnson & Zhang (2013) as follows:
m = ff(x, 8 V fxx &,) + (F))),
AGI Arwna
where X is an anchoring point (a.k.a,, reference point) that updates periodically. This variance-reduced stochastic gra- dient can reduce the variance of the stochastic gradient by
Notations In this paper, we assume x, denotes the parameter of the language model at step and E1...., Er E = are a
2

---

MARS: Unleashing the Power of Variance Reduction for Training Large Models
adding a correction term (x) + VF(X) based on less frequently updated reference point x and its full gra- dient VF()). It can be shown that the variance my can be controlled by ||x; - x| which will diminish as both x, and & converges to the stationary points when the algorithm makes progress, Subsequent improvements in variance re- duction techniques were introduced in SARAH (Nguyen al., 2017a) and SPIDER (Fang et al., 2018), which get rid of the anchor point and result in the following momentum update:
3 Method
In this section, we introduce MARS (Make vAriance Reduction Shine), family of preconditioned optimization algorithms that perform variance reduction in gradient esti- mation.
3.1 MARS Framework
We first introduce our framework for a preconditioned. variance-reduced stochastic optimization, which unifies both first-order (e.g.. AdamW, Lion) and second-order (e.g.. Shampoo) adaptive gradient methods.
m. = ff(,, - (x(xx-1, + m
Preconditioned Variance Reduction. Variance reduction methods achieve faster convergence than SGD, yet identi- rying optimal learning rates remains a practical challenge. Particularly, different parameters often exhibit varying cur- vatures, requiring tailored learning rates for cach. One ap- proach to addressing this issue is to use the Hessian matrix to precondition gradient updates, integrating curvature infor- mation into the updates. The idea stems from minimizing the second-order Taylor expansion at Xx:
x = X m.
(2.3)
In the context of training neural networks, x; E Rd repre- sents the trained weights in the neural network, & repre- sents random data, and m, is the variance-reduced (VR) first-order momentum. The stochastic gradient difference term (x(,,, V Ex) cancels out common noise brought by E., while pushing the gradient estimation from the estimator of V FF (x to the estimator of FX(x).. full gradient (or a large batch stochastic gradient) Fx(x),, which we refer to as anchoring step, analogous to the anchor point in SVRG.
However, m, needs to be reset periodically
F(x++1) = F(x)) + VF((x) xx))
2 (x - x₁) F(x)(x)(xx+)) - xx),, (3.1)
Subsequently, Cutkosky & Orabona (2019) introduced Stochastic Recursive Momentum (STORM), variant of standard momentum with an additional term, achieving the same convergence rate as SPIDER while eliminating the need for periodic anchoring:
resulting in the update formula X X; - Fxx)).
where H, == F(x()) E Raxd is the Hessian matrix. In our paper, we encapsulate the preconditioned gradient H, F update within a more generalized framework of Online Mirror Descent (OMD) as in Gupta et al. (2018), leading to the following update rules:
m, -1 (1 B1) V1)
{
31 (V Ex) vf(x-1,
(2.4)
X,41 arg min XER
T) (m,, x) + - x|--
2 H.
(3.2)
where 8 0 is momentum parameter, and (1(V((x {r {x)) is the additional term that has variance reduction effect. Note that if x, x X;-1, STORM becomes approximately the standard momentum. Alternatvely, (2.4) can be rewritten as an exponential mov- ing average (EMA) of the first order momentum from previ- ous step/iteration and the stochastic gradient with a gradient correction term:
where 1): > 0 can be viewed as a base learning rate. Com- bining (3.2) with the STORM momentum, we obtain the following preconditioned variance-reduced update:
m, -3mm-1- (1- (1-3:) B1) [ fxx, E:)
B1
+ 1- B1
(V (Vf(xx, Ex) - v ff(x kil)
(3.3)
X++1 arg min XER
T)t (m, x) + 2 ||xx - xx |
(3.4)
m = 3,m
+ (1
- 31) V Ex),
31
V (Xe, - -x(xx1,,))))
(2.5)
Remark 3.1. SuperAdam (Huang et al., 2021) also incor- porates the STORM into the design of adaptive gradient methods. However, their precondition matrix can be viewed special case of our general framework. Super design focuses on diagonal precondition matrix and draws heavily from the design used in Adam (Kingma & Ba, 2015), AdaGrad-Noma (Ward et al., 2020), and AdaBelief (Zhuang et al., 2020). Furthermore, their preconditioner matrix is designed following Adam's 's structure but does not account
B
gradient correction
Theoretically, when assuming access to an unbiased stochas- tic first-order oracle to the objective function F(x). STORM achieves the nearly optimal gradient complexity of O(E-3) for non-convex and smooth optimization problems (Arjevani et al., 2023).
3

---

MARS: Unleashing the Power of Variance Reduction for Training Large Models
for the revised definition of variance-reduced momentum, resulting in a significant mismatch between the first-order and second-order momentum. We will further clarify these differences when discussing specific instances of our frame- work.
estimator:
E X [[X] -
= Var(X) 2y [[ - (Y - -[[x] + + Var(Y),
which admits an optimal choice of y that minimizes the variance:
Algorithm Design. In practice, alongside our precondi- tioned variance-reduced update (3.3), we introduce a scaling parameter ?/ to control the scale gradient correction in variance reduction. We also introduce new gradient esti- mator C₁ which is the combination of stochastic gradient and the scaled gradient correction term:
arg min E [X E[X] XX E[Y]
- - E[X EYDDI
Var(Y)
B1 (Vf(x,, Et) (x(
C, = VF=v((xx fx(( +
In the context of STORM updates used in our work, the update rule includes both stochastic gradients and recursive momentum terms. Let us define:
sealed gradiens coreccion
When 7 1, the above reduces the second term of (3.3). On the other hand, when 7/- 0, it reduces to the stochastic gradient. Thus, can be seen gradient estimator with adjustable variance control.
X (1 ((x)()))
Y 8 [ E(+1) (x(xx,
2. m, Vx(x))
Following standard techniques in deep learning practice, we also perform gradient clipping on C₁ which is calculated by:
and let U == X - + Z, The updated Zt+1 at step [++ is of the form + (+Y - E[Y]), whose squared expectation we aim to minimize. The optimal choice of y in this setting is:
C
if |c > 1,
C = Clip(c,, 1)
=
Ct
Taill
otherwise.
(3.5)
EU + Var(Y)
E[Y
We note that the Second-order Clipped Stochastic Opti- mization (Sophia) algorithm (Liu et al., 2023) also incor- porates clipping in their algorithm design. However, their approach does clipping upon the preconditioned gradient hipping-byy-vluee by-value, while our method applies clipping with clipping-
With this choice, we obtain variance reduction:
EIU + - E I - (Y -
(E[UY) + Var(Y))²
to the intermediate gradient estimate using the more stan- After the gradient
E|Y3
dard technique of Hopping-by-none clipping, the VR momentum m, can be calculated as the EMA of C, The resulting MARS algorithm summarized in Algorithm 1.
which is strictly smaller than the variance under any non- optimal choice of Thus, dynamically tuning % improves the variance of the gradient estimator, leading to better con- vergence behavior. A full convergence analysis is provided in Appendix B.2.
Algorithm MARS
1: input: Xo, 31, {rt). {
We provide the convergence analysis of Algorithm 1 in Theorem B.5 in Appendix B.2. We prove that under standard assumptions, MARS achieves a superior convergence rate of O(T-1/3), outperforming the OTT4)/4 rate attainable by AdamW.
2: Set mo + 0 and X1 + X0
3: for = 1, to n do
4: Sample &. and let Ct = =Vf(x,, (x(x,,
3. 7.T (V/(x,, Ex) V (x)
5: if [f|ccl12 > 1, then o; = e. clse C2 = C2
Full Matrix Approximation. In practice, calculating the Hessian matrix is computationally expensive or even in- tractable due to the complexity of second-order differenti- ation and the significant memory cost of storing H,, espe- cially when the parameters in neural network constitute high-dimensional matrix. Many existing algorithms em- ploy various approximations of the Hessian. For instance, K-FAC (Martens & Grosse, 2015) and Shampoo (Gupta et al., 2018) approximate the Gauss-Newton component of the Hessian (also known as the Fisher information matrix),
6: m. m, + (1 31)C;
7: x = arg min [n, (m, x) + -|11 X 1
8: end for
Why %/ improves convergence A similar idea of adjusting the strength of variance reduction has been proposed by Yin et al. (2023) in the context of SVRG. This approach originates from a classical line of work on control vari- ates (Asmussen & Glynn, 2007; Lavenberg et al., 1977). In the standard control variates setting, one considers the

---

MARS: Unleashing the Power of Variance Reduction for Training Large Models
using a layerwise Kronecker product approximation (Mor- wani et al., 2024). Additionally, Sophia (Liu al., 2023) suggests using Hutchinson's 's estimator the Gauss-Newton- Barlett estimator for approximating the Hessian. We take various designs of the preconditioning matrix into account and broaden the definition of H, in (3.4) to encompass vari- ous specifically designed preconditioning matrix in the rest of the paper.
3.2.1 MARS-ADAMW
The first instance of MARS is built up on the idea of Adam/AdamW (Loshchilov & Hutter, 2019). To automat- ically adjust the learning rate and accelerate convergence, Adam (Kingma & Ba, 2015) adopts the adaptive precon- ditioned gradient in (3.6) together with a bias correction and lz regularization. A (Loshchilov & Hutter, 2019) further changes the lg regularization decoupled weight decay. Overall, the full AdamW updates can be summarized as follows:
Diagonal Matrix Approximation. Even when using ap- proximated Hessian matrices, second-order algorithms men- tioned above remain more computationally intensive com- pared to first-order gradient updates. Thus, another line of research focuses approximating the Hessian matrix through diagonal matrices, as seen in optimization algo- rithms like AdaGrad (Duchi et al., 2011), RMSProp (Tiele- man, 2012), AdaDelta (Zeiler, 2012), Adam (Kingma & Ba, 2015) and AdamW (Loshchilov & Hutter, 2019), etc. This approach to diagonal preconditioning effectively transforms the updates into first-order method, assigning adaptive learning rates to each gradient coordinate. For example, in AdaGrad (Duchi et al., 201 1), the preconditioned matrix is defined by:
m: = 3 8,m=- + (1 [x)
(3.7)
V, B2 (1 (1)-)2)(V/))))) f fx E:))
(3.8)
m
Vi
m. 1 - 31
V. 1 - 32
(3.9)
m + Xxx
Xt+1 Xt TJ
(3.10)
Vi +E
We see that except for the small e introduced for com- putational stability, and the decoupled weight decay Ax, AdamW can be seen as step of mirror descent update (3.2) with m, defined in (3.7), v, defined in (3.8), and H, defined by
H
[x)(
3 35
H == V diag vi)
(3.11)
V1
On the other hand, Adam can be seen using a diagonal Hy, where each diagonal element is the EMA of [x](x,, E:)];
In MARS-AdamW, we implement the preconditioned variance-reduced update in (3.4), and utilize the same definitions for H,, 6, and weight decay as those specified in AdamW. For vt, different from the EMA of squared gradients m in Ad we redefine it to fit our variance- reduced stochastic gradient. Specifically, we denote the summation of the stochastic gradient and the scaled gradient correction term by c, and define v, as the EMA of of as follows:
H HlH +11 (1 (1-1)]Vf(xx))) fxx, E.).
(3.6)
Therefore, the update simplifies to elementwise adaptive gradient update, i.e., [x = [xx]; 77 ///mml/// / Our unified framework accommodates both types of precondi- tioning: full Hessian approximation and diagonal Hessian approximation. Different definitions of H, give rise to dif- ferent algorithms.
C = V.((x, Et)
B1
Notably, full-matrix approximations of the Hessian are po- tentially more powerful than diagonal approximations, they capture statistical correlations between the gra- dients of different parameters. Geometrically, full-matrix approximations allow both scaling and rotation of gradients, whereas diagonal matrices are limited to scaling alone.
+ 7t 1 31 (Vf(x,, Ei) (((x
(3.12)
m. 3,m,-1 + (1 -31)0ct
(3.13)
V: = 32V1-1 + (1 - (1-B2)c3 B2)c²
(3.14)
Here, 71 is a scaling parameter that controls the strength of gradient correction. When % = 0, the algorithm reduces to AdamW. Conversely, when ? = 1, (3.13) aligns with the STORM momentum. Combining (3.12), (3.13), (3.14) together with (3.11) and the mirror descent update (3.2), we derive the MARS-AdamW algorithm in Algorithm 2. In practice, is often set between 0 and 1. Moreover, we employ gradient clipping- sopping--y-nome by-norm to at Line 5, following the standard gradient clipping technique performed in neural network training. We provide a convergence analysis of A Algorithm 2 in Theorem B.6 in Appendix B.2.
3.2 Instantiation of MARS
In previous subsection, we introduced our preconditioned variance reduction framework in Algorithm 1 and discussed various approaches for approximating the Hessian matrix. In this subsection, we introduce practical designs of MARS un- der different choices of H. While here we only present three instantiations: MARS-AdamW. MARS-Lion, and MARS-Shamppe we believe there are many other instances of MARS can be derived similarly.
5

---

MARS: Unleashing the Power of Variance Reduction for Training Large Models
Remark 3.2. Compared with SuperAdam (Huang et al., 2021), one key difference is that our algorithm defines the second-order momentum Vt as the exponential moving aver- age of the square norm of C2 rather than the square norm of the stochastic gradient. This new definition of second-order momentum is crucial for accommodating the right scale of updates on a coordinate-wise basis. Moreover, as we mentioned in Algorithm 1, we introduce scaling parameter T and implement gradient clipping on C₁- In Section 4, we will demonstrate empirically that the changes contribute to effective performance in large language model training. Finally, our algorithm utilizes bias correction and weight decay while SuperAdam does not.
Algorithm 2 MARS-AdamW
1: input: Xg, X, $1, S2. } {nc}
2: Set my < 0, Vo + 0 and X1 f- Xu
3: for 1. = 1, to n do
Sample & and let Ct P1
v)(((, +
1-8 (Vf Et) - vf(x)-1)
if cf > 1, then č₁ Tillla else C₁ = Ct
6:
m, = f,mm- + (1- (1-B))
7:
Ve = 82V1-1 + (1 ---ba -
m = T- Vi = vi
m
-
X X2 1) (Vitee
m.
+ Xx,
10: end for
Remark 3.3. Careful readers might have noticed that in each iteration of our algorithm, we need to calculate the stochastic gradient twice for different data batches &: and & with the same parameters. In order to overcome this problem, we propose to use (Vf(x, E,) - -V((xx1: to approximate (Vf(x,,&)) - (f(xx,,1, in (3.12) and C will be approximated by:
which employs sign operation to maintain uniform magni- tude across all parameters. The updates for Lion are illus- trated as follows:
m, = :2u, 1 + (1 (1- f Et),
(3.15)
u, -1u, + (1 -1) (x,, Et),
(3.16)
C, = Ex),
Bi " 11- 31
(Vf(x,, Ex) (x
X++1 = Xx T) sign(m:) + Xxx)
(3.17)
To avoid confusion, we refer to the approximate version as MARS-approx. While MARS and MARS-approx dif- fer in their updates and may theoretically exhibit distinct convergence guarantees, our experiments show that MARS provides only marginal improvements over MARS-approx in practice. Thus, we recommend using MARS-approx for practical applications.
Instead of employing an EMA of gradient norms as in (3.8) and (3.11) of AdamW, the sign preconditioning mechanism in Lion utilizes
H / diag(m?)
(3.18)
Following the same definition of H, as in (3.18), we present MARS-Lion in Algorithm 3.
Connection between MARS-Adam' and Adan. Adan (Xie et al., 2024) is another adaptive gradient method improved upon Adam with reformulated Nesterov's accelerated SGD (See Lemma 1 in Xie et al. (2024) for more details). The Adan algorithm takes the following momentum updates:
Algorithm 3 MARS-Lion
1: input: xo, X, $r, {Te), {nd
2: Set my <- 0 and x1 f- X0
3: for 1. = 1, to do
4: Sample &. and let C2
vf(x), +
f 't T- (V/(x,, E:) - (((x
y: = 31 + (1 - -8:) (x),);
5: if | > 1, then ट₁ न.ी else č₁ Ct
z, = 3,2,- (1 32) (V)(((( &(xx))
m. = m,m; + (1 --000
X x, 7) (s.gn(m:) + xxxx)
m; : yt + 82Zz.
8: end for
When 32 = B1, this reduces
Setting m, = V fxt. Et), a1 = 31 02 = 1 B2, by = B2. b2 = 1- - 33 in Lemma D.1, and shifting the index of u, by taking u, = u, u = u in Lemma D.I, we can show that Lion momentum updates in (3.15) and (3.16) are equivalent to the following single momentum update:
m, -3 + (1 3.) Et) V (xx-111
+ +
which is special case of MARS-approx's momentum with Nt 1- 31. It is worth noting that although motivated by the Nesterov's momentum, Adan's 's momentum updates cannot recover Nesterov's 's momentum unless 31 = =B2-
m, = 8 9mm,-1 + (1 31) (xx, Et)
(3.19)
+ B1 (2)
1 1- 1-31 (x)() (((x(11
E(-))
3.2.2 MARS-LION
(3.20)
Using symbolic program search. Chen et al. (2023) in- troduced a simpler algorithm Lion compared to AdamW,
On the other hand, setting ?t = -192 in the core updates
6

---

MARS: Unleashing the Power of Variance Reduction for Training Large Models
of MARS (3.12) and (3.13), we obtain the m, update:
clipping at Line 3, we can see Lion a special case of MARS-Lion when " 1 1 and using approximate gra- dient calculation on V fx,,1,, ..). In practice, we observe little difference between using f1 derived from the STORM momentum and its approximation f(x ((
m, = f, + + (1- 81)
V Ex)
(3.21)
31 32)
+
(V (Vf(x, Et) (1 {.
1 1- B1
(3.22)
3.2.3 HARS-SHAMPOO
The only difference between (3.26) and (3.28) lies in the stochasticity used, specifically, & versus &--1 when cal- culating (= Therefore, ignoring the gradient clipping at Line 3, we can see Lion as special case of MARS-Lion when ,/2 = 11- and using approximate gra- practice, we observe little difference between using f(x-11, Et-1) derived from the STORM momentum and its approximation (,
Shampoo (Gupta et al., 2018) introduces preconditioning approach that operates on the eigenspace of matrices. Given the gradient matrix m, == (x,(x, Ex) E Rmxx the update rules of Shampoo are displayed follows:
L, = L + m,,m,
dient calculation on (f(x
R. RR-1 + m, mr,
X(+1 =X,- MeLL 1/4 R,
1/4 (3.29)
Connection between MARS-Lion and Lion. Lion turms out to be a special case of MARS-Lion. The momentum updates in Lion can be seen as an approximate implemen- tation of our updates. To facilitate this claim, we present lemma that follows directly from straightforward arithmetic calculations.
where X, Rmxx (slightly abusing notation) represents the corresponding weight matrix. It has been shown that the two-sided preconditioning in (3.29) is equivalent to pre- conditioning on the flattened vector m, == vec(m;) with Kronecker product (Gupta et al., 2018; Morwani al., 2024)
Lemma 3.4. For any sequence fgr E Rd )==,,,.... con- sider the following updates of m, for any constant factors a1, 02 b,, and by:
H [ (G.,)) & (Σ (E.ra.)"
Ur aj 1 + a2gr.
(3.23)
m; = b u: + b2g.-
(3.24)
In practice, an exponential moving average (EMA) often used in place of the direct summation. The up- date rule in (3.29) can be simplified to X X,
The updates are equivalent to
T) (m. -1/1 m, (m, J m.) -1/1 performing preconditioning on the eigenspace of m,
This is equivalent to
m, = a-mm + (bla2 - anbb + b2)g; + a1b2(g. g g(-1).
Setting g: = /x Et). a. = 31, 02 = 11-31.66 B1 bi = 32. b2 = - 332 Lemma D.1, and shifting the index of u, by taking u, 4- u u in Lemma D.I. we can show that Lion momentum updates in (3.15) and (3.16) are equivalent to the following single momentum update:
U, Σx, V. = =SVD(m:)
X X, n7 U.V.
(3.30)
Therefore, we borrow the eigenspace preconditioning from Shampoo, and design our algorithm to precondition on any matrix-shaped update as in (3.30). In particular, we present our algorithm in Algorithm 4.
m, = 3 mmm:11 + +(1 (1- 31)
(3.25)
91 (1) +
1 - B1 (Vf(xe, Ex) - (1
Algorithm 4 MARS- MARS-Shampor Shampoo
1: input: Xo, X,, X, f1, {T)), {nd)
(3.26)
2: Set mn <- 0 and +- Xo
On the other hand, setting 7/- = 1 32 in the core updates of MARS (3.12) and (3.13), we obtain the m, update:
3: for t = 1, to n do
4: sample &. and let Ct = -V((x),)) +
8: fV &.) - (f(((-1, 7e(15)
m, = 8,m,-1 1 + (1 (1) Ex)
(3.27)
5:
m, = - (1
6:
U, V, SVD(m))
+ 31 (2)
(V Ee) (f(xx-1) +1)
7:
X = X Th (U.VT + Axx)
B1
end for
(3.28)
To reduce the time complexity of SVD decomposition, Bern- stein & Newhouse (2024) summarized four different ap- proaches for computing or approximating (3.30) includ- ing SVD, sketching (Martinsson Tropp, 2020), Newton
The only difference between (3.26) and (3.28) lies in the stochasticity used, specifically, & versus ε when cal- culating vf Therefore, ignoring the gradient
7

---

MARS: Unleashing the Power of Variance Reduction for Training Large Models
iteration (Lakić, 1998; Higham. 2008; Anil et al., 2020), and Newton-Schulz iteration (Schulz, 1933; Higham, 2008).
4.1 Experimental Setup
All our experiments done based the nanoGPT (Karpa- thy, 2022) implementation of the GPT-2 (Radford et al., 2019) architecture, and on the OpenWebTex (Gokaslan et al., 2019) dataset. The training and validation sets contain approximately 9 billion and 4.4 million tokens, respectively, all preprocessed using the GPT-2 tokenizer. We conduct experiments on three scales of GPT-2 models: small (125M parameters), medium (355M parameters), and large (770M parameters). Per the nanoGPT configurations, we disabled biases, applied GeLU activations, and set the Dropout rate (Srivastava et al., 2014) to 0.0. We utilized 16 NVIDIA A100 GPUs for training the small models. For the medium and large models, training was conducted on 32 NVIDIA A100 GPUs and 32 NVIDIA H100 GPUs, respectively. Other hyper-parameters of training are listed in Appendix F,
Our algorithm design accommodates any of these SVD solvers to best fit specific computational needs.
Connection between MARS-Shampoo and Muon. Muon (Jordan et al., 2024) is recently proposed algo- rithm that utilizes the Newton-Schulz iteration (Higham, 2008; Schulz, 1933) to solve the SVD problem. It has demonstrated superior performance in terms of convergence speed when compared with AdamW and Shampoo train- ing large language models. The update rules of Muon are demonstrated as follows:
u, 1 + Et),
(3.31)
m, la + Ex);
(3.32)
O, = NewtonSchulz (m.))
X1+1 X - (O, + )x,).
4.2 Results
Applying Lemma D.1 to (3.31) and (3.32), with m ff(x,, a1 = 14, 02 = 1, bj = p. b2 = 1. we obtain an equivalent single update of momentum:
In Figure and Figures 2-3 (in the Appendix), we demon- strate the training and validation losses as a function of training tokens and wall-clock time for various model sizes² Across the small, medium, and large GPT-2 models, MARS consistently surpasses both the AdamW and Muon base- lines in training and validation losses. The performance gap becomes more pronounced with increasing model size. No- tably, MARS exhibits both rapid initial decay and sustained superiority throughout the training process. Further, we ex- plore the performance of additional learning rate choices in Appendix E. Notably, the best validation losses of MARS- AdamW and MARS-Lion achieved in our GPT-2 large ex- periments are 2.511 and 2.534. For comparison, the best validation losses are 2.568, 2.565 and 2.606 for AdamW, Lion and Muon. These results demonstrate that our reported performance is highly competitive with state-of-the-ar opti- mizers.
m 10 V
+ (Vf Ei - V ((-))))
(3.33)
On the other hand, taking 31 u., 7t = 1 1= 1- in MARS, (3.12) and (3.13) reduces to
m.
pmm-1 + - VX)
+ + - () Ex)
(()
By dividing both sides of the above equation by 1 - 44, we obtain
m.
m + E.)
1
1 + (x( (V (3.34)
In Figure 1(c), well as Figures 2(c) and 3(c) in the Ap- pendix, we compare the wall-clock time of different al- gorithms. We observe that MARS-AdamW and MARS- Lion have a slightly higher per-iteration cost compared to Adam W but is much faster than Muon. Additionally, they consistently demonstrate lower validation losses than both AdamW, Muon and Lion within equivalent training dura- tions.
In can be seen that (3.34) is a rescaled version of (3.33), except that the stochastic gradients V fx Ex) and f Ex) are taken both at &
Experiments
In this section, we evaluate the performances of two instanti- ations of our algorithm, MARS- MARS-AdamW Adamm and MARS-Lion
in comparison with Adam W (Loshchilov & Hutter, 2019), the predominant algorithm for training large language mod- els, Lion (Chen et al., 2023) and Muon (Jordan et al., 2024) on GPT-2 model series. More experiment results and abal- tion study, including the computer vision experiments, the effect of different learning rate schedulers, as well as sensi- tivity to y and batch size, are postponed to Section E.
We also evaluate O-shot and 5-shot performances of our optimizer on common benchmarks including ARC (Ya- dav et al., 2019). BoolQ (Clark et al., 2019), Hel- laSwag (Zellers et al., 2019), OBQA (Mihaylov et al., 2018), PIQA (Bisk et al., 2020), WinoGrande (Sakaguchi et al., 2020) and MMLU (Hendrycks et al., 2021), with the im-evaluati i oon on-harness codebase (Gao et al., 2024). We only list the 5-shot performances for large mod-
For the sake of training efficiency, we use MARS- aaprox for the experiments as the default configuration, except for Appen- dices E.2 and E.4. Discussion of the difference in performance between MARS-exace and MARS -ppprox is in Appendix E.2.
'The training loss curves are smoothed using Exponential Mov- ing Average.
8

---

MARS: Unleashing the Power of Variance Reduction for Training Large Models
GPT-2 Large (770M). openwebtext
GPT-2 Large (770M), openwebtext
3.
5.
GPT-2 GPT-2 (770M), upenieebeest
AdamW/
AdamW
Lion
Lion Muon
AdamnW
Nuon
MAlts-ddemi
MA-8- MM
Munn
MART-AddemiV MARS Uso
Toeliing Laa
2.8-
Vallcatin Loss
Massaction uass
2.7-
2.2
2.5
2:7-
2.5
2.5
2.6-
20000
40000
60000
80000
100000
Thiining
Traning
Waleelcck time on 32XH100 (saconds)
(a) Training Loss
(b) Validation Loss
(c) Wall-clock time
Figure 1. The training and validation loss curves, plotted against both training tokens and wall-clock time GPT-2 large model (770M).
Table 1. The evaluation results of large models pre-trained using the OpenWehTex Text dataset (5-shot with sarpaluuatio-hhameests
The best
scores each column bolded. Abbreviations: HellaSwag HellaSwag, WG WinoGrande.
METHOD
ARC-E
ARC-C
BOolQ
HELLASWAG
OBQA
PIQA
WG
MMLU
Sc1Q
AVG.
ADAMW
52.95
28.67
56.33
42.55
29.40
67.68
52.01
25.27
82.90
48.64
LION
52.53
26.88
52.42
43.41
29.80
67.63
54.46
24.70
85.70
48.61
MUON
49.58
26.88
55.78
40.42
30.20
66.65
$2.64
24.58
79.10
47.31
MARS-ADAMW
54.04
26.28
62.78
45.66
31.60
68.12
$2.49
25.93
84.50
50.15
MARS-LION 54.25
28.92
56.36
44.00
29.20
69.10
54.93
25.98
85.50
49.80
els in Table and leave other results in the Appendix E.I. The models pre-trained with HARS-AdamW and MARS- Lion outperform those pre-trained with AdamW, Muon and Lion optimizers, validating an enhanced downstream perfor- mance within the same number of pre-training steps.
Impact Statement
This paper presents work whose goal is to advance the field of optimization theory in deep learning. We believe that our work contributes meaningfully to the field, specifi- cally on advancing the efficiency in the pre-training stage of deep learning models, especially Large Language Models. By involvement of variance reduction in adaptive learning methods, our method can greatly lower the cost for pre- training language models on limited training corpora and more resource-constraingg
5 Conclusion
In this work, introduce MARS, unified framework for adaptive gradient methods that integrates variance reduction techniques to improve the training of large models. Our approach combines the adaptive learning rate introduced by preconditioning with the faster convergence enabled by variance reduction. Within our framework, we have devel- oped three optimization algorithms based on the ideas of AdamW, Lion, and Shampoo. Through extensive empirical experiments on GPT-2 pre-training tasks, we demonstrate that MARS consistently outperforms baseline algorithms in terms of both token efficiency and wall-clock time. Our re- sults establish a generic framework for combining adaptive gradient methods with variance reduction techniques, con- tributing to the advancement of optimizers in large model training.
devices well as in broader settings, opening new avenues for their application in vari- ous downstream tasks. Improvement in efficiency typically correlates with reduced energy consumption, potentially de- creasing the environmental footprint of L.LM pre-training. This advancement underscores the potential of optimization method development in deep learning field in both techno- logical and societal contexts.
References
Allen-Zhu, Z. and Yuan, Y. Improved svrg for non-strongly- objectives. In International conference on machine learning, pp. 1080-1089. PMLR, 2016.
convex or
Anil, R., Gupta, V., Koren, T., Regan, K., and Singer, Scalable second order optimization for deep learning. arXiv preprint arXiv:2002.0901 2020.
Arjevani, Y., Carmon, Y., Duchi, J. C., Foster, D. J., Srebro, N., and Woodworth, B. Lower bounds for non-convex

---

MARS: Unleashing the Power of Variance Reduction for Training Large Models
stochastic optimization. Mathematical (1):165-214, 2023.
AaooammaacaaPPorammmmettt Programming, 199
Cutkosky, A. and Orabona, F. Momentum- variance
reduction in non-convex sgd. Advances in neural infor- mation processing systems, 32, 2019.
Asmussen, S. and Glynn, P. W. Stochastic simulation: algo- rithms and analysis, volume 57. Springer, 2007.
Defazio, A. and Bottou, L. On the ineffectiveness of vari- ance reduced optimization for deep learning, Advances in Neural Information Processing Systems, 32, 2019.
Bernstein, and Newhouse, L. Old optimizer, new norm: An anthology. arXiv preprini 1rXiv:2409.203222 2024.
Defazio, A., Bach, F,, and Lacoste-Julien, S. Saga: A fast incremental gradient method with support for non- strongly convex composite objectives. Advances in neural information processing systems, 27, 2014.
Bisk, Y., Zellers, R., Bras, R. L., Gao, J., and Choi, Y. PIQA: reasoning about physical commonsense in natural language. In The Thirty-Fourth AAAI Conference on Arti- ficial Intelligence, AAAI 2020, The Thirty-Secont Inno- vative Applications of Artificial Intelligence Conference, IAAI 2020, The Tenth AAAI Symposium on Educational Advances in Artificial Intelligence, EAAI 2020, New York, NY. USA, February 7-12, 2020, pp. 7432-7439. AAAI Press, 2020.
Defazio, A., Yang, X. A., Mehta, H., Mishchenko, K., Khaled, A., and Cutkosky, A. The road less scheduled. arXiv preprint 1.Xx::2405.166822 2024.
Derezinski, M. Stochastic seriance-reduced newton: Ac- celerating finite-sum minimization with large batches. In OPT 2023: Optimization for Machine Learning.
Brown, T. B. Language models are few-shot learners. arXiv preprint arXiv:2005. 14165, 2020.
Devlin, Bert:
Pre-training of deep bidirectional trans- formers for language understanding. arXiv preprint arXiv: 04400, 2018.
Carlson, D., Cevher, V., and Carin, L. Stochastic spectral descent for restricted boltzmann machines. In Artificial Intelligence and Statistics, pp. 111-119. PMLR, 2015a.
Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Yang, A., Fan, A., al. The llama 3 herd of models. arXiv preprint arXiv:2407.21788 2024.
Carlson, D., Hsieh, Y.-P., Collins, E., Carin, L., and Cevher, V. Stochastic spectral descent for discrete graphical mod- els. JEEE Journal of Selected Topics in Signal Processing, 10(2):296-31,, 2015b.
Chen, J., Zhou, D., Tang, Y., Yang. Z., Cao, Y., and Gu, Q. Closing the generalization gap of adaptive gradient methods in training deep neural networks. arXiv preprint arXiv: 1806 1806.06763, 2018
Duchi, J., Hazan, E., and Singer, Y. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 2011.
Fang, C., Li, C. J., Lin, Z., and Zhang, T. Spider: Near- optimal non-convex optimization via stochastic path- integrated differential estimator. Advances in neural in- formation processing systems, 31, 2018.
Chen, X., Liang, C., Huang, D., Real, E., Wang, K., Pham, H., Dong. X., Luong, T., Hsieh, C.-J., Lu, Y., et al. Sym- bolic discovery of optimization algorithms. Advances in neural information processing systems, 36, 2023.
Frangella, Z., Rathore, P., Zhao, S., and Udell, M. Promise: Preconditioned stochastic optimization methods by incor- porating scalable curvature estimates. Journal of Machine Learning Research. 25(346):1 2024.
Chowdhery, A., Narang, S., Devlin, J., Bosma, M., Mishra, G., Roberts, A., Barham, P., Chung, H. W., Sutton, C., Gehrmann, S., et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240):1-113 2023.
Gao, L., Tow, J., Abbasi, B., Biderman, S., Black, S., DiPofi,
A., Foster, C., Golding, L.,, Hsu, J., Le Noac h, A., Li,
Clark, C., Lee, K., Chang, M., Kwiatkowski, T., Collins, M., and Toutanova, K. Boolq: Exploring the surprising difficulty of natural yes/no questions. Burstein, J.. Do- ran, C., and Solorio, T. (eds.), Proceedings of the 2019 Conference of the North American Chapter of the Associ- ation for Computational Linguistics: Human Language Technologies, NAACL-HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume (Long and Short Papers), pp- 2924-2936. Association for Computational Linguistics, 2019.
H., McDonell, K., Muennighoff, N., Ociepa, C.,, Phang,
J., Reynolds, L., Schoelkopf. H., Skowron, A., Sutawika,
L., Tang, E., Thite, A., Wang, B., Wang, K., and Zou, A. A framework for few-shot language model evaluation,
07 2024. URL https: /ttp///eeooooooo/ooo wwweoooooooooooooooooooooooooooe records/ 12608602.
Ghadimi, S. and Lan, G. Stochastic first-and zeroth- methods for nonconvex stochastic programming. SIAM journal on optimization, 23(4):2341-2368, 2013.
10

---

MARS: Unleashing the Power of Variance Reduction for Training Large Models
Gokaslan, A., Cohen, V., Pavlick, E., and Tellex, S. Open- webtext corpus. http://Skyl1 Mockyylon077.oom tiub.. io/ 1o/Openn@ebTextCorppe
Kavis, A., Skoulakis, S., Antonakopoulos, K., Dadi, L. T. and Cevher, V. Adaptive stochastic variance reduction for non-convex finite-sum minimization. Advances in Neural Information Processing Systems, 35:23524 -2538, 2022.
2019.
Graves, A. and Graves, A. Long short-term memory. Super- vised sequence labelling with recurrent neural networks, pp. 37-45, 2012.
Kingma, D. P. and Ba, Adam: A method for stochastic optimization. In Bengio, Y. and LeCun, Y. (eds.), 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Confer- ence Track Proceedings, 2015.
Gupta, V. Koren, T., and Singer, Y. Shampoo: Pre- conditioned stochastic tensor optimization. Interna- fional Conference on Machine Learning, pp. 1842-1850. PMLR, 2018.
Krizhevsky, A., Hinton, G., et al. Learning multiple layers of features from tiny images. 2009.
Hazan, E.,, Levy, K., and Shalev-Shwartz, S. Beyond con- vexity: Stochastic quasi-convex optimization. Advances in neural information processing systems, 28, 2015.
Lakić, S. On the computation of the matrix k-th root. ZAMM-Journal of Applied Mathematics and Mechan- ics/Zeitschrift für Angewandte Mathematik und Mechanik: Applied Mathematics and Mechanics, 78()):677172,, 1998.
He, K., Zhang, X., Ren, S., and Sun, J. Deep residual learn- ing for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Lavenberg. S. S., Moeller, T. L., and Welch, P. D. The application of control variables to the simulation of closed queueing networks. In Proceedings of the 9th conference on Winter simulation- Volume 1, pp. 152-154, 1977.
Hendrycks, D., Burns, C., Basart, S., Zou, A., Mazeika, M., Song, D., and Steinhardt, J. Measuring massive multitask language understanding. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021, 2021.
LeCun, Y., Botton, I.., Bengio, Y., and Haffner, P. Gradient- based learning applied to document recognition. Proceed- ings of the IEEE, 86(1 1):2278-2324, 1998.
Higham, N. J. Functions of Matrices. Society for Industrial and Applied Mathematics, 2008
Levy, K., Kavis, A., and Cevher, V. Storm+: Fully adaptive sgd with recursive momentum for nonconvex optimiza- tion. Advances in Neural Information Processing Systems, 34:20571 2021.
Hu, S., Tu, Y.. Han, X., He. C.,, Cui, G., Long. X., Zheng. Z., Fang, Y., Huang, Y., Zhao, W., et al. Minicpm: Unveiling the potential of small language models with scalable train- ing strategies. arXiv preprint urXiv:2404.06395, 2024.
Li, H. Smoothness and Adaptivity Nonlinear Optimiza- tion for Machine Learning Applications. PhD thesis, Massachusetts Institute of Technology, 2024.
Huang. F., Li, J., and Huang, H. Super-adam: faster and universal framework of adaptive gradients. Advances in Neural Information Processing Systems, 34:9074-9085, 2021.
Li, H., Rakhlin, A., and Jadbahaie, A. Convergence of adam under relaxed assumptions. Advances in Neural Information Processing Systems, 36. 2024.
Johnson, R. and Zhang, T. Accelerating stochastic gradient descent using predictive variance reduction. Advances in neural information processing systems, 26, 2013.
Liu, A., Feng, B., Wang, B., Wang, B., Liu, B., Zhao, C., Dengr, C., Ruan, C., Dai, D... Guo, D., et al. Deepseck-v2: A strong, economical, and efficient mixture-of-expperts language model. arXiv preprint arXiv:2405.0444
2024.
Jordan, K., Jin, Y., Boza, V., Jiacheng, Y., Cecista, F., New- house, I.., and Bernstein, J. Muon: An optimizer for hidden layers in neural networks, 2024. URL https: likelleejordann haaerdanngittttt github. ic wttptsts////ooooooccoccottsss /muon/.
Liu, H., Li, Z.,, Hall, D., Liang, P., and Ma, T. Sophia: A scalable stochastic second-order optimizer for language model pre-training. arXiv preprint 2rXiv:2305.14344 2023.
Kaddour, J., Key, O., Nawrot, P., Minervini, P., and Kusner, M. J. No train no gain: Revisiting efficient training algo- rithms for transformer-based language models. Advances in Neural Information Processing Systems, 36, 2024.
Liu, M., Zhang. W., Orabona, F., and Yang. T. Adam + A stochastic method with adaptive variance reduction. arXiv preprint arXiv:201 1. 1985, 2020.
Karpathy, A. NanoGPT. https:: //tp:/bggtttt////// com/ ka karpathy/nanoo 2022.
Loshchilov, and Hutter, Decoupled weight decay regu- larization, In 7th International Conference on Learning

---

MARS: Unleashing the Power of Variance Reduction for Training Large Models
Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019, 2019.
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I., al. Language models unsupervised multitask learners. OpenAl blog, 1(8)99, 2019.
Lozhkov, A., Ben Allal, L., von Werra, L., and Wolf, T. Fineweb-edu: the finest collection of educational content, 2024, URL beattps:/huugiin ace.. co/ datasets htwp://ngaacee//Eieeeeee edu.
Reddi, J., Hefny, A., Sra, S., Poczos, B., and Smola, Stochastic variance reduction for nonconvex optimization. In International conference machine learning, pp. 3 314- 323. PMLR, 2016.
Martens, J. and Grosse, R. Optimizing neural networks with kronecker- factore approximate curvature. In Interna- tional conference machine learning, pp. PMLR, 2015.
Reddi, S. J., Kale, S., and Kumar, On the convergence of adam and beyond. arXiv preprint arXiv: (XXiv::904.0023331 2019a.
Martinsson, P.-G. and Tropp. J. A. Randomized numer- ical linear algebra: Foundations and algorithms. Acta Numerica, 29:403-572, 2020.
Reddi, J., Kale, S., and Kumar, S. On the convergence of adam and beyond. arXiv preprint arXiv: 1904. 1904.092377 2019b.
Riedmiller, M. and Braun, H. A direct adaptive method for faster backpropagation learning: The rprop algorithm. In IEEE international conference on neural networks, pp. 586-591. IEEE, 1993.
McCandlish, S., Kaplan, J., Amodei, D., and Team, O. D. An empirical model of large-batch training. arXiv preprint arXiv: arXiv:1812.06162. 2018.
McMahan, H. B. and Streeter, M. Adaptive bound opti- mization for online convex optimization. arXiv preprint arXiv: 2010.
Roux, N., Schmidt, M., and Bach, F. A stochastic gradient method with an exponential convergence .rate for finite training sets. Advances in neural information processing systems, 25, 2012.
Mihaylov, T., Clark, P., Khot, T., and Sabharwal, A. Can suit of armor conduct electricity? A new dataset for open book question answering. In Riloff, E., Chiang, D., Hock- enmaier, J., and Tsujii, J. (eds.), Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, Brussels, Belgium, October 31 November 4, 2018, pp. 2381-2391. Association for Computational Linguistics, 2018. doi: 1.. I 1260. URL https: /ttp:://do...o orgg /0. 1ww1111661lli/ 8- 1260.
Sakaguchi, K., Bras, R. L., Bhagavatula, C., and Choi, Y. Winogrande: An adversarial winograd schema challenge at scale. In The Thirty- AAAI Conference on Arti- ficial Intelligence, AAAI 2020, The Thirty- Second Inno- vative Applications of Artificial Intelligence Conference, IAAI 2020, The Tenth AAAI Symposium on Educational Advances in Artificial Intelligence, EAAI 2020, New York, NY, USA, February 7-12, 2020, pp. 8732-8740 AAAI Press, 2020.
Morwani, D., Shapira, I., Vyas, N., Malach, E., Kakade, S., and Janson, L. A new perspective on shampoo's preconditioner. arXiv preprint arXiv:2406.17748 2024.
Saon, G., Kurata, G., Sercu, T., Audhkhasi, K., Thomas, S., Dimitriadis, D., Cui, X., Ramabhadran, B., Picheny, M., Lim, L.-L,, et al. English conversational telephone speech recognition by humans and machines. arXiv preprint arXiv: 173.02/703.22336 2017.
Nesterov, A method for solving the convex program- ming problem with convergence rate o(1/k²). Proceed- ings of the USSR Academy of Sciences, 069:543-547, 1983. URL https: http//ppp..aaasseeeeeeeee www.nticschollr.......cc.c TTrp/CoopusID:14511.5
Schölkopf, B. and Smola, A. J. Learning with kernels: support vector machines, regularization, optimization, and beyond. MIT press, 2002.
Nesterov, Y. Introductory lectures on convex optimization: A basic course, volume 87. Springer Science & Business Media, 2013.
Schulz, G. Iterative berechnung der reziproken matrix. Angew. Math. Mech,, 13:57-59, 1933.
Nguyen, L. M., Liu, J., Scheinberg, K., and Takáč, M. Sarah: A novel method for machine learning problems using stochastic recursive gradient. In International conference on machine learning, pp. 2613-2621. PMLR, 2017a
Shalev- Shwartz, S. and Zhang, T. Stochastic dual coordi- nate ascent methods for regularized loss minimization. Journal of Machine Learning Research, 14(1), 2013.
Shazeer, N. and Stern, M. Adafactor: Adaptive learning rates with sublinear memory cost. International Con- ference on Machine Learning. pp. 2018.
Nguyen, L. M., Liu, J., Scheinberg, K., and Takáč, M. Stochastic recursive gradient algorithm for nonconvex optimization. arXiv preprint 1705v: 2017b.
PMLR,
12

---

MARS: Unleashing the Power of Variance Reduction for Training Large Models
Shi, H.-J. M., Lee, T.-H., Iwasaki, S., Gallego-Posada, J., Li, Z.,, Rangadurai, K., Mudigere, D., and Rabhat, M. A distributed data-parallel pytorch implementation of the distributed shampoo optimizer for training neural net- works at-scale. arXiv preprint arXiv:2309. 2023.
Zeiler, M. D. Adadelta: adaptive learning rate method. arXiv preprint arXiv: /212.500, 2012.
Zellers, R., Holtzman, A., Bisk, Y., Farhadi, A., and Choi, Y. Hellaswag: Can machine really finish your sentence? In Korhonen, A., Traum, D. R., and Mārquez, L. (eds.). Proceedings of the 57th Conference of the Association for Computational Linguistics, ACL 2019, Florence, Italy, July 28- August 2, 2019, Volume 1: Long Papers, PP- 4791-4800. Association for Computational Linguistics, 2019. doi: (1.18653/V1/P199 1472. URL https:// doi. org/ /00 1tt../1165/v/lppp 1472.
Srivastava, N., Hinton, G., Krizhevsky, A., Sutskever, I., and Salakhutdinov, R. Dropout: simple way prevent neural networks from overfitting. The journal of machine learning research, 15(1):1929-1955 2014.
Tieleman, T. Lecture 6.5-rmsprop: Divide the gradient by running average of its recent magnitude. COURSERA: Neural networks for machine learning, 4(2):26, 2012.
Zhang, S., Roller, S., Goyal, N., Artetxe, M., Chen, M., Chen, S., Dewan, C., Diab, M., Li, X., Lin, X. V., et al. Opt: Open pre-trained transformer language models. arXiv preprint arXiv:2205.01663 2022a.
Vaswani, A. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
Vyas, N., Morwani, D., Zhao, R., Shapira, I., Brandfon- brener, D., Janson, L., and Kakade, Soap: Improving and stabilizing shampoo using adam. arXiv preprint arXiv:2409. 11321, 2024.
Zhang, Y., Chen, C., Shi, N., Sun, R., and Luo, Z.-Q. Adam can converge without any modification on update rules. Advances in neural information processing systems, 35: 28386-28399, 2022b.
Wang, Z., Ji, K., Zhou, Y.. Liang, Y., and Tarokh, V. Spi- derboost and momentum: Faster variance reduction al- gorithms. Advances in Neural Information Processing Systems, 32, 2019.
Zhao, R., Morwani, D., Brandfonbrener, D., Vyas, N., and Kakade, S. Deconstructing what makes a good optimizer for language models. arXiv preprint Niti:2200.079722 2024.
Ward, R., Wu, X., and Bottou, L. Adagrad stepsizes: Sharp convergence over nonconvex landscapes. Journal of Ma- chine Learning Research, 21(219):1-30, 2020.
Zhou, D., Xu, P., and Gu, Q. Stochastic nested variance re- duction for nonconvex optimization. Journal of machine learning research, 21 21(103): 1-633, 2020.
Xie, X., Zhou, P., Li, H., Lin, Z., and Yan, S. Adan: Adap- tive nesterov momentum algorithm for faster optimizing deep models. IFEE Transactions on Puttern Analysis and Machine Intelligence, 2024.
Zhou, P., Xie, X., Lin, Z., and Yan, S. Towards understand- ing convergence and generalization of adamw. IEEE Transactions on Pattern Analysis and Machine Intelli- gence, 2024.
Yadav, V., Bethard, S., and Surdeanu, M. Quick and (not so) dirty: Unsupervised selection of justification sentences for multi-hop question answering. In Inui, K., Jiang, J., Ng. V., and Wan, X. (eds.), Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing, EMNLP-IJCNLE 2019, Hong Kong, China, November 3-7, 2019, pp. 2578-2589. Association for Computational Linguistics, 2019. doi: 10. 10.118653V//DDI9-1D00 URL https:://// /tt.. N0.18653//11D19--20
Zhuang. J., Tang, T., Ding, Y., Tatikonda, S. C., Dvornkk, N., Papademetris, X., and Duncan, J. Adabelief optimizer: Adapting stepsizes by the belief in observed gradients. Advances in neural information processing systems, 33: 18795-18806, 2020.
Yin, Y., Xu, Li, Z., Darrell, T., and Liu, Z. coefficient makes svrg effective. arXiv preprint 1Xiv:2311.005888 2023.
You, Y., Li, J., Reddi, S., Hseu, J., Kumar, S., Bhojanapalli, S., Song, X., Demmel, J., Keutzer, K., and Hsieh, C.-J. Large batch optimization for deep learning: Training bert in 76 minutes. arXiv preprint arXiv: /004.0990.0000 00960, 2019.
13

---

MARS: Unleashing the Power of Variance Reduction for Training Large Models
Appendix
A Related Work
15
Theoretical Analysis
15
B.1 Connection to Nesterov's Acceleration
15
B.2 Convergence of MARS
16
Proof of Theorems
17
C.1 Proof of Theorem B.5
18
C.2 Proof of Theorem B.6
19
Proof of Auxiliary Lemmas
21
D.1 Lemma D.1 and Proof
21
D.2 Lemma D.2 and Proof
22
D.3 Proof of Lemma C.I
22
D.4 Proof of Lemma C.2
23
D.5 Proof of Lemma C.3 and C.4
25
D.6 Proof of Lemma C.5
26
Additional Experiment Results
26
E.1
Supplementary Results for the Main Experiments
26
E.2
MARS and MARS-appros
26
E.3
Experiments on FineWeb-Edu 100B Dataset
27
E,4
Computer Vision Experiments
28
E.5
Sensitivity to 7.
30
E,6
Different Learning Rate Scheduler
31
E.7
Sensitivity Batch Size
Hyper-parameter Settings
34
14

---

MARS: Unleashing the Power of Variance Reduction for Training Large Models
A
Related Work
In this section, we provide review of additional related works, including some previously mentioned, to help readers gain deeper understanding of the history and development of adaptive gradient methods and variance reduction techniques. Adaptive Gradient Methods. RProp (Riedmiller & Braun, 1993) is probably one of the earliest adaptive gradient methods by dynamically adjusting the learning rate. AdaGrad (Duchi et al., 2011; McMahan & Streeter, 2010) adjusts the learning rate based on the geometry of the training data observed during earlier iterations. To tackle with the issue of diminishing gradient in AdaGrad (Carlson et al., 20:5a; Tieleman (2012) introduced RMSProp by incorporating the idea of exponential moving average. A significant advancement came with Adam (Kingma & Ba, 2015), which integrated RMSProp with Nesterov's momentum (Nesterov, 1983; 2013) achieving superior performance and becoming a prevalent optimizer in deep neural network training- Later, Loshchilov & Hutter (2019) proposed to decouple weight decay from gradient calculations in Adam and introduced AdamW, an optimization algorithm having become the predominant optimization algorithm in contemporary deep learning applications. To fix the convergence issue of Adam, Reddi al. (2019b) introduced the AMSGrad optimizer, which maintains running maximum of past second-ordes momentum terms achieve non- non-increasing step sizes. Subsequently, Chen et al. (2018) unified AMSGrad and SGD within the Padam framework by introducing a partial adaptive parameter to control the degree of adaptiveness. Notably, AdamW and its variations have been widely used in the training of popular large language models, including OPT (Zhang et al., 2022a), Llama (Dubey et al., 2024), and DeepSeek- V2 (Liu et al., 2024).
Variance Reduction Methods. SAG(Roux et al., 2012) and SDCA(Shalev-Shwaal & Zhang, 2013) were among the first attempts to apply variance reduction techniques to accelerate the convergence of SGD. Subsequently, simpler algorithms like SVRG(Johnson & Zhang, 2013) and SAGA(Defazio et al., 2014) were introduced, achieving the same improved convergence rates. SARAH (Nguyen et al., 2017al further simplified these approaches by employing biased recursive gradient estimation, which reduces storage requirements while achieving the complexity bounds for convex optimization problems. And some researchers have also attempted apply preconditioning into variance reduction in the convex setting (Frangella et al., 2024; Derezinski). For non-convex optimization, besides SVRG (Allen-Zhu & Yuan, 2016; Reddi et al., 2016) and SARAH (Nguyen et al., 2017b), SPIDER (Fang et al., 2018) integrates Normalized Gradient Descent (Nesterov, 2013; Hazan et al., 2015) with recursive estimation of gradients, while SNVRG (Zhou et al., 2020) introduces multiple reference points for semi-stochastic gradient calculation for improved variance reduction and convergence rate. SpiderBoost (Wang et al., 2019) refines SPIDER by enabling the use of significantly larger constant step size while preserving the same near-optimal oracle complexity. Subsequently, STORM (Cutkosky & Orabona. 2019) was proposed to further simplifies the SPIDER and SNVRG algorithms through the use of stochastic recursive momentum. This was later improved into a parameter-free variant, namely STORM+(Levy et al., 2021).
Variance Reduction for Adaptive Gradient Methods. Few works have explored the application of variance reduction techniques to adaptive gradient methods. To the best of our knowledge, the only exceptions are Adam SuperAdam and AdaSPIDER, Adam+ (Liu et al., 2020) attempts to reduce the variance of first-order moment into Adam by estimating the gradient only at extrapolated points. SuperAdam (Huang et al., 2021) and VRAdam (Li, 2024) integrates variance reduction with AdamW to achieve improved convergence rates. And AdaSPIDER (Kavis et al., 2022) introduced adaptive step size in SPIDER algorithm. However, these variance-reducee adaptive gradient methods have primarily been validated on basic computer vision tasks, such as MNIST (Schölkopf Smola, 2002) and CIFAR- (Krizhevsky et al., 2009), and simple natural language modeling tasks, like SWB-300 (Saon et al., 2017), using straightforward architectures such as LeNet (LeCun al., 1998), ResNet-32 (He al., 2016). 2-layer LSTMs (Graves & Graves, 2012), and 2-layer Transformers (Vaswani, 2017). As result, significant gap remains in the successful application of variance reduction techniques to adaptive gradient methods, particularly in the rapidly evolving domain of large language models.
B Theoretical Analysis
B.1 Connection to Nesterov's Acceleration
Many optimization algorithms exhibit similarities with Nesterov's acceleration and can be considered adaptations of Nesterov's s acceleration with varying parameterization schedules, as discussed in works such as Defazio et al. (2024) and Xie et al. (2024). In this section, we compare and contrast Nesterov's momentum and STORM momentum used in our paper. Specifically, Nesterov's accelerated gradient descent can be equivalently written (Xie al., 2024):
m; = B m + + (/
15

---

MARS: Unleashing the Power of Variance Reduction for Training Large Models
while the STORM momentum is
m. 81 m + (1 B:) (i)V 31
( (VI E Vaa-:-)---((-----.. (x(x)-,,
The most significant difference lies in the noise handling schemes. In Nesterov's acceleration, Vx((;, E.) is subtracted by V fx (t-1) to determine a direction of improvement. In contrast, STORM variance reduction subtracts V (.) from f to cancel out the noise introduced by &.. Furthermore, in our theoretical analysis (Section B.2), we prove that variance-reduced variants of Adam W achieve an improved convergence rate of (TT-1/3). Empirically, we also show that a variance reduced noise schedule performs better than its approximate counterpart.
B.2 Convergence of MARS
Although Kingma & Ba (2015) did convergence analysis for Adam, Reddi et al. (2019a) pointed out that they made some mistakes in the proof, and they also proved that in some special cases, Adam does not converge. However, there are some attempts to prove the convergence of Adam and AdamW in special circumstances (Zhang et al., 2022b; Li et al., 2024; Zhou et al., 2024). Our algorithm, MARS, is also based on Adam W. In addition, it involves the property of variance reduction. We prove that MARS can converge with a better convergence rate with careful selection of hyperparameters. To help analyze the convergence of the algorithm,
make the following assumptions:
Assumption B.1 (Bounded Variance). We assume that the variance of gradient estimator is bounded by 0². i.e., for any noise &. parameter and fF()) = [<, {x), there exists positive such that:
E (x, [ S So²
(B.1)
Assumption B.2 (I-Smoothness)) We assume that for arbitrary& &. fax E) is L-smooth:
f(x. -flr - L||x - y1l2, yx
(B.2)
Assumption B.3 (H Lower Bounded). We assume that there is a constant p > 0 such that for all H2, > 0, H, > pL. Remark B.4. Note that this is implicitly satisfied by the instantiations of MARS since we add a small to semi-positive definite H; for computational stability.
We proposed Theorem B.5 for our main Algorithm and Theorem B.6 for MARS-AdamW (Algorithm 2, where additional weight decay is involved). We note that for theoretical analysis, it is necessary to consider vire-na parameters B1, and 32,r. However, in practice, these parameters are typically set constants.
Theorem B.5. In Algorithm 1, under Assumptions B.1, B.2 and B.3, when choosing 7) = (s + 1)-333 2 80 Suppose e 2 321.2p 2 + 1, BLa41 =1 1 - ont and $2,
1 - v/² then VT 2 it holds that
E
m1113 5 (20G (xx,)
pe-
log(s +T))
1
p² 2==1 M+11
T
4L2
T2/3
81/T1/33
1
E - x1112 5
16G
+
20²0²
log(s (og(s+T))
1
ΣL M411
T
m
3p
3L2
T2/3
6T.771/3
where G = Fxi)) - miny F(x) +
162 and Mo+1 is defined in (C.2).
Theorem B.6. In Algorithm 2, under Assumptions B.1, B.2 and B.3, when choosing The (s + 1)-1/3, >
max(8L3/P)) /p²,64x3) 64X²). Suppose ||x|l2 I D, 32L2p=² 2 + 1, 91.41 = =1 - on and B2 = 1 - m/, then VT 2 s, it
holds that
T E Vxx:))
m2 (2p(C+ + DOOGG>DD2oogg Log(s + T)) + pe302
loggs + +T))
1
P² ET M411
4L2
T2/3
16L2T1/3
16(G + AD log(s log(s +T))
20202
loggs + +))
1
E1=1 M++1
+ 1
1/t
P
T2/3
1/T1/3
where G= F(x1) - minx +(x)+ (2 +) +
13 161
and M is defined in (C.2).
16

---

MARS: Unleashing the Power of Variance Reduction for Training Large Models
The theorems above guarantee the convergence of ( We remark that even though it seems that the involvement of weight decay may results in slower convergence, it performs better in practice, la Loshchilov & Hutter (2019). Finally, the term Mc+1 is always greater or equal to 0, and when 7/+11 I, we would have M++1 0. The dependency of Mc+1 on 7++1 is shown in (C.2) and Lemma C.2. This proves that the convergence is faster if we allow flexible "t schedule.
C Proof of Theorems
First, we introduce auxiliary lemmas necessary for proving the theorems.
Lemma C.I. In Algorithm 2, for any 0 S $2,, I 1 and VI 2 1, the following inequality holds:
Vlv. V & V2211 (2,1)
(C.1)
The proof is in Section D.3. Lemma C.2. In Algorithm Under Assumption B.I and B.2, if 2 B1,+11 2 0, Vt, under approximate choice of 7+1 as in (D.12), we have
E| VF(()))))))))))))))))))))))))))) m S BVF((.) mel3 2.8 1 T.2
x.12 + 2(1 -1,,41)002 M
where
A2+1 (31 -1-1111-++1)) - Act1)) ) == EIV VX(,
(C.2)
Ge+1 + (Voentr(VVVV/((xx11..55 - v[(xx,
A41:= =
Vf E++1) Vf(xx
and
G411 = (1- B1.+1)BB Vf(x0+1, Vf(xx,, Vf(xx+1 VF(x(+1)
+ 31,411 (F(x++) VF(x)), F(x:) m
The proof is in Section D.4.
Lemma C.3. In Algorithm 1. With Assuptions B.2 and B.3 and n) S p. (2L) -1 -1, Vt 2 1, holds that
F(x+11) S F(x:)
P 27/
-1xx X
+ Th
|VF(x,) m
p
The proof of Lemma C.3 in Section D.5.
Lemma C.4. In Algorithm 2. With Assuptions B.2 and B.3 and ne & min .p: (2L)-1), "), 2 1. it holds that
X F(x++1) + xf4 H441844 F(x:) + A
2 x, H.x; P. ||x; - x
477.
IVF(xx) - mml1? + 2 2(1 D2.
The proof of Lemma C.4 is in Section D.5.
Lemma C.5. Let 7)+ = (s + -) 1/3 a 2 1, VI 2 0. Then th2 nr S nr Vt 2
17

---

MARS: Unleashing the Power of Variance Reduction for Training Large Models
C.1 Proof of Theorem B.5
Proof of Theorem B.5. First, we define the Lyapunov function
i = E F(x))
+
P 16 /IVF((x) m,
vt 2 1.
Then we calculate the difference between two consecutive Lyapunov functions as:
B+ 4; = =E(F(x(+1) F(x,)
+ E
P
|VF((x++) - m 11
P
VF(xx) - -m
(C.3)
16L2M)
161-1
12
For 11, we use Lemma C.3 to obtain
P. ||x: -
p
241112+ "h
|[VF(x;) - -ma12
(C.4)
217
P
For I2, we use Lemma C.2 to obtain
1 = E
P
[[VF(x(-))
m
P
[F m mil1
16L2m
16L2
1
p&R.et11
p(l - -1,,4413002
(x)x)
mell
El|x+++1 - xe112
M
16T2
116
1
8m)
81.21t
161.21
(
82
1
P
p=2
-
EE | x1
Mr+1,
(C.5)
1612
7
T/-
St/e
8L2
16L2n
where the last inequality follows from the definition that Blet1 = 1 - cni Further, for the first term on the right hand side, we have
P
B,,41
1 S 16L2
1 1
(Butt)
1 -
P
(1 ctf
17 17-1 1
P
(
em
16L2
T
T/-1
n
TT-1
16L2
1612
17-1
From Lemma C.5, we know that
Mh.
Choosing e such that > 32L p obtain
-t-1
B7,41
1
16L2 (n -mn) < -2ppp
(C.6)
16L
N)
M--1
Bringing (C.6) into (C.5), we arrive at the upper bound for Iz:
1 5 2n V
mel1
P
EE
X1+1
xel13
pc211002
P
M141
(C.7)
p
St/e
8L2
16L2n
Now combining (C.3), (C.4) and (C.7), derive
$ p++1 o S Tl E|| [x))) - my
3p
X1 B||x++1
El|
1/3
p
mp
M
8m/
8I.2
161.21
Taking a telescoping sum for t 1,... T gives
E VF(x))
mll1 -
3p
E E xx s Φ Tr+1 +
pc
1
T
p
M
8m
8L2
+
16L2m
==1
1=1
2
Φ 1 +
log(s +T)
p
M
812
16L2m
By the definition of d., we have P7++ 2 F(XT+1) 2 miny F(x). And for d..
1/3
1/3
ps2/302
Φ = E. F((((x))
[VF(x)) -mi
= Fx(x)
167,2 B[IVF((x)) - < F(x1) +
161.
161.2
18

---

MARS:
Unleashing the Power of Variance Reduction for Training Large Models
0² 161
Consequently, defining G = F(xi) - minx F(x) +
the following inequality holds:
T
T
VF(x))
mel1 +
3p
X x
G +
pc² 822T
log(s +
Σ
T
M411
(C.8)
8m,
T
T 16L2n
1=1
Dealing with the two terms the left hand side separately, have
T
PG
p²60²
p²
V
mill x
+
log(s +T)
M
T
Tur
8L2TMI
T t=l
16L2
pG + p²620
(8 + 7)2/3
loggs + 77)
M++1
81.2
T
16I
pe20²
p² ΣL= 1t= M+11
2pG +
log(s + T))
4L2
T2/3
812T1/3
where the last inequality holds when 2 8. Similarly, the second term (C.8), we also have
1
3p
G
pc² 8L2T
P PE M4+1
E B
x
log(s T)
T
81z
16L272/8
which implies that when T 2 8, due to the definition 1).
8G
202
Σt M411
E x
x1112 -
log(s +
T
m
3pI'm
3L2TAT
6T272233
16G
26
Σt= M4+1
log(s + +T)
3/T2/3
312T22/3
6T,2T1/3
That concludes our proof.
C.2 Proof of Theorem B.6
Proof of Theorem 8.6. First, we define the Lyapunov function as
A +. (xx) + 2 X, H,x: + 161.m--1 VF(xx) m. 12]
14 2 1.
Then we calculate the difference between two consecutive Lyapunov functions as:
A X, H2x,
F(x++1) + x He.H F(x.)
+ E
m 112 [xF(()++)
p ||VF(x;) m.
(C.9)
161.2
161.2m-1
12
For Ij, we use Lemma C.4 to obtain
XD2 /2(1 B22)
P.
||x:
1x41112
1
Tt. |VF(x;)
m
(C.10)
47
19

---

MARS: Unleashing the Power of Variance Reduction for Training Large Models
For 12, we use Lemma C.2 to obtain
12 =
E
m
P
[V(xx)
m
l1
1612TM
1612
(111-1
1
18| (xx))
m
p&R.ett
Bl|x++1 - xx112 -x
p(1 - 11,2+1)2022
P
Mc+1
1612
1]t-1
8t)t
81211/
161,1/
P
[R..-1
B
(x)x))
m
+
P .E x - x
pe nj.
M4+1,
(C.11)
1672
Th-1
8t}e
SL2
1612
where the last inequality follows from the definition that 31,141 = 1 - eni- Further, for the first term on the right hand side, we have
B7,41 T/
1
B1.411
1
1 - cni
1
p
(=
1
1
cTH)
1672
1/-1
16T2
T/
T--1
161,2
1/t
1/-
1672
1)t
T/-1
From Lemma C.5, we know that
Choosing such that 2 32I22p
obtain
n
P
1
(n cmt) S -2m
(C.12)
16L
T
Th--1
16L2
Bringing (C.12) into (C.11), we arrive at the upper bound for I2:
1 S 2/k =|V - m +
P
pc211002
EE
X
x
M1+1-
(C.13)
P
8t)
8L2
16L21
Now combining Formulas (C.9), (C.10) and (C.13), obtain
P+11 01 S Th E(((x)))
m 1
P
E
x - xe12
pc 02
AD2
/
2/111
32,1)
P
M
8m)
8L2
2
16L2m
Taking telescoping sum for t 1,... T gives
U VF(())
mz +
-B|x1111 -xx- x1113)
1=1 P
8m)
pc T
XD2 12(111(1- - (2,1)
T
P M,41
< 1 r411
+
8L2
8 +
16L2
T
I 11 4TT1 +
pc²
log(s s TT) + XD2 log(s + Σ
P
M441,
8L2
16L2m
where the last inequality follows by taking 82, = 1 - - By the definition of 4,, we have 4T+1 2 F(x+1) 2 miny F(x).
And for 1 according to the fact following (D.7) that | [iag( + c)
- + c, we obtain
Φ 1 (E
X x2
p/2/3
H +
m.
167
< F(x)) + A (2((1+c) c)+
ps-/3
E((x)(((,) -
v1(x),
2
16I2
< F(x)) + > 2
(2(1+ c)+ ps
16I2
Consequently, defining G = F(xi) - min F(x) + A D2 + + + 1/3
1612 the following inequality holds:
T
1
M VFF(x) m, 112 + P.
E E|Exx X1+1 xx13)
(C.14)
T
8mp
1=1
G
2c202
log(s TT) + 2D loglog(s +T) 1
<
M
(C.15)
T
8L2T
16L2Tp
20

---

MARS: Unleashing the Power of Variance Reduction for Training Large Models
Dealing with the two terms on the left hand side separately, have
1
mll p(G + log + T))
p2202
1 =1=1 M4+1
[x))(())
-log(s + T)
T
TAp
8L
16L271/3
pC + AD log(s +T)) + p²6=02
(s + T)
p² D M
log(s log(s+T)) +T)
8L2
T
161,211/3
(2p(G+ AD2 log(s +T) + pc²
p²
EE= M++1
log(s 22)
412
T2/3
16I2TI/3
where the last inequality holds when 2 s. Similarly, for the second term, we also have
T
p
El
Xx+1
x1112 G XD2 log(s +T)
+
pc²
log(s + +
2 M
8T
8I.T
161.2Tp
which implies that when T 2 due to the definition of 1/+
B||x++11 -xx111xxe11--- 8(G + XD2 log(s + T))
200²
ΣT Mc-1
B
llog(s + T))
T
z f
pTTnr n/
T.TTT
2T.2T
16(G + XD2 log(s + +T)
202
Ei= Met1
log(s + +7)
2/2
22T2/3
1277/3
That finishes the proof.
D Proof of Auxiliary Lemmas
D.1 Lemma D.1 and Proof
Lemma D.I. For any sequence (m; E ) consider the following updates of m, for any constant factors 01 02 b1, and b2:
u, dj u, a2gr.
(D.1)
= b + bagi.
(D.2)
The updates are equivalent to
m, = 0-m--11 + (b102 0rd2 + b2)g. + a.bb(gg -(-1).
Proof of Lemma D.I. Substituting (D.1) into (D.2), we obtain
m. by (a) u + aggi) + begi = aubiu,111 (bla2 + by)g.
(D.3)
On the other hand, shifting the index of (D.2) by 1, it holds that
m,-1 = b b2g.-1-
(D.4)
Combining (D.3) and (D.4), we obtain the iterative update of m, from its previous value m,-1 as:
m. aj m a1b2g.-1 + (bia2 + b2)g.
a m - (bi02 a1b2 + b2) + albz(g. - -g--)).
This completes the proof.
21

---

MARS: Unleashing the Power of Variance Reduction for Training Large Models
D.2 Lemma D.2 and Proof
Lemma D.2. In Algorithm 2, assume there is constant D such that |xx|la D for all t 0. Given that 0 I 82, S
for all ( 0, the following inequality holds:
p(1 - 7)
A x Hr+1X; H - - -2.1182.)DD))
(me, X2 X(+1)
1Ixe
X
(D.5)
1/t
2
Proof of Lemma D.2. By definition of m, and the update rule of Xx+1 in Algorithm 2, we have
1
(m;, X - X141)
H.((1 - -x x++1), X1 x(+1)
1/1
p(1
meA)
|x1 x (Hi X x(+1),
(D.6)
1/1
where the inequality follows from Assumption B.3. By convexity of h1(x)) == f T Hix, we obtain
1 x H1x+1
H2x - x++1) 1 xx H2X
2
Therefore, (D.6) becomes
(me, X1 - - X(+1) 2 p(1 - -)
H2x, + A
Xx X++1
122
2 x
2 2x1+ Hix
p(1 1)
À X H(41)x--
[x x
1/2
2
x
H2,, + 2 H. 1 +
T)
We recall that H; = diag( Vv + <) in Algorithm 2. Combining this with Lemma C.I. we derive
x HH, H4t1)x441 x (v. V
2 2(1 - D2.
Overall, we conclude
(m,,x, x(+1) p( m:)) 17c
A xx, H H,x -2 B2,
|xc x 11
D.3 Proof of Lemma C.1
Proof of Lemma C.I. According Algorithm 2, is the clipped c, with the norm |/ 1. Therefore, bound v.
by:
|vll2
Σ 32,
II B2g j=k+1
II B2, vo
Σ 82,1)
II j=
B2.
(1
II $2, 1
(D.7)
where the first inequality is due to vo = O, and the second inequality holds since 0 S 82, S 1. We note that when ts = t, we treat 12 $2.j as 1. Similarly, since mo = 0 and 0 < 31,t I 1, we have upper bound of m, as:
|
Σ 11 k=1
Cx
II BI.
B mo
(1 31,6)
II Bij = (1 II g=b+1
B
1..
(D.8)
j=
1=1
k=1
k=1
Therefore, according the Ve+1 update in Algorithm 2, we have
|| - Velloo
|[1 - 32,) (0: vr)lloo
< (1 (( + l|velloo)
< (1 (2,1)(( ||vel22
s 2(1 12,e),
22

---

MARS: Unleashing the Power of Variance Reduction for Training Large Models
where the first inequality is due triangle inequality and the second inequality derives from that |x]|oc IV vyl S vI= yl. Vv, 2 holds that
||x||2- Since
IIvv Vetillo S VIlv - S 22 32,1)
D.4 Proof of Lemma C.2
Proof of Lemma C.2. By the definition of m, in Algorithm 1,
11.141 E(11) -V (xxe
m441 91, + (1 31, (f(x E(+1) + 7++1 1- 31,141
-1- V + 31,141 me 7 (V (V(()))))))
Subtracting both sides by (x obtain
m1+1 VF(x++1)
= (1 - 31,41 ( + B1,111 m. 744 (VI(x()), - (i((x, VF(x++1)
= -11,e41 [m - (x())) + 31.4 VF(()) VFF(x++)
- -1.141 V + V vf(x(+1:(44)) - Vf(xi: Setz))
= -B.tt1 - VF(()) +1 B1,041) (f( (AA
- VE(X++I)
+ B1.41 (VF(x))-VF(Xx))) VF(X + +44181444111 -
Rearranging the terms, we get
m++1 (F((++)==++++)))))++)))))))))))) = B1.41) (1((x++1:4+1)) -7/0x4114HH)VV/((x))) (F(H)))))((((+)))) + 11.0+1(1m, VF(xi))
+ B1,411
VF(xx) - FF(x++1 + 71+1 ((4))) V f Ee+1))
(((x)
With shorthand of notations, we write m: becomes
VF(x:), and A. (x(()+1,,441) Vx(x,, The above
Ex+1 = (1=
VF(x++1) + B1,441E, (7,44(70111 EAt);
(D.9)
where the expectation in the last term taken over the randomness in E,41. Taking squared norm over both sides of (D.9) and then take expectation over Ex+1, we have
E e1 12 E1[(1 31,++1 (V E(t1) VF(x+11) B1,411, B1,111 (141 (14411111444 EA,) 11
E1I(1 31,++1
- VF((++1) B1,411 + B (A EA, - (7+11 - 1)
E][(1 - B1,e+1) (VI( - VF((41)) + BL.411, + B1,41 A -EA,) 1
32 -74441114411 1)
1I
2(1 - 31.13 E
VF(x++1)) A., (1 (Vixs))((/((x))) -VF(x)) +
81,411; + 11,411(A - EAA)
111
(D.10)
23

---

MARS: Unleashing the Power of Variance Reduction for Training Large Models
For I, observe that E is independent of É++1 and the expectations of (X(+1114441) V Fx((++1) and A, EA,
all 0. Therefore
= ==Bl[ -
- VF + 31,+1 (A: - EA:) 112 + 1 e
2 f1,++1 E V ( VF(x-1111 + 2 f, E A EA, 12 B 1
(D.11)
Minimizing 11 11-2(1 7 (1+=1),,,4+1111 over 7141 in (D.10), we know that when
A., (1 - (V VF((x+1)) + B1,411 + 31,41 EAA))
E
1-141
B 11
II - 1I-2(1 70+1)81,0-1111 III reaches optimality at
B A., (1
VF(x++1) + S1.tte+ (I, BA.)) (EA-))))
11 - 2(1 11+1)81.+1111===
1.1.
Using G. to represent
G-- (1 -1-==1111,1111 B1 A.. E(1) - VF(x++1) (
and defining
Gi+1 + $1,141 (E EA, B)
A
We have
2
II 2(1 7441) (1 III
El A (1,44(11- -7+11 Act1
1EIA,11242.
and when
GG+1 + -1,2-1 (B A 3 - EEA,
(IEA,) 12 G++1
7+11 1
(D.12)
31,4+111 B A
BI, EE
11 - 2(1 - 31. 111 -BIA,,1242.
Defining
M
= :=
E 3 BIA. (1.t1(11- - 7141) - Ast2
we conclude that
min
11
tr..e.....
< 1| M
< 2(1 - 31,++ / [[ E--1) VF ) 2,8 E| A. EA, + Ble,11 -M+11
< 2(1 ) 2 + 289 _ 1|x1111
xx113
S E| M
where in the last inequality we utilized the fact that the terms finishes the proof.
A,
E A
3 & E A 113 and the L-smoothness of f. Rearranging
24

---

MARS: Unleashing the Power of Variance Reduction for Training Large Models
D.5 Proof of Lemma C.3 and C.4
Proof of Lemma C.3. Given the L-smoothness of F(x) in Assuption B.2, we have
F(x(11) I F(x:) + (VF( X++1 xx) + 2 ||x(+1 - xcl1 F(x:) + (m,. Xt+1
L xx) + (VF(x)) m. X1+1 x1) 2
||x+1 - X1
(D.13)
By definition of m, and the update rule of X in Algorithm we have
(m, X X(+1)
H. [x[x
X441l, X(+1)
1/1
P
|xx: x1-1
(D.14)
1/1
Here the inequality follows from Assumption B.3.
Proof of Lemma C.4. Bringing (D.14) into (D.13), we obtain
F(x++1) S F(x:)
P |xt X++1 112 + (VF(x)) m., X++1 - -xx) + 2 ||x++11 x:
L
1/8
< (xx:)
P
||x X 112 U
VF(x;) me + P
|xx+1 x1l12 + 2/2 -||x++1 x
?/1
4rje
S F(x))
P 1xce
x
VF(x,) 71t p
m.l12.
2T
The second inequality follows from applying both the Stuchy-Schward and Young's inequalities. Moreover, the final inequality results from selecting 7, satisfy n # This completes our proof. Bringing (D.5) Lemma D.2 into (D.13), we have
F(x(11)= < F(x)) + (m, X++1 - xi) + FV)(x) - m. Xx+1 xx) + L
||x-11 1
< Fxx,)= + (VF())) - m, X - Xx) + -|x - xell3
p(1- n)
|x X 12
2 x, Het1xe - x Hexe+1 2(1 - D21.)
1/:
Taking t)- < min miin((4) ((41)-1
11 pp- (2T.) we have
p(1 mc))
|x x++1 113
L
||x++1 - xx (-
3p
+
4
||x
xxoll2 <
| -xel22
778
2
41/
277.
Therefore,
P
F(x++1) $ F(x:) + (VF(x)) - m. X,+1 - xii)
||x: - x1+1
21/1
x H x H 12 -32
2
By Cauchy-Schwarz inequality and Young's inequality. we have
(V)():) m,, X xx) [x,)) mell2 -||x+11 xx|12 Ttt [VF(x))
mm.12 112 +
||x x+112.
4nj.
Therefore, we conclude that
F(x++1) & F(x.) + Th [VF(x)) m2
P | - xxl12
P ||x X, x
41p
217:
X 1 | H H 1 H H2
/2(1 D22
X [x] H,,1X, H +
A 2 V D2
F(x,) Vx,)) m, 11
||x+1
-x+
47/-
2
Rearranging terms finishes the proof.
25

---

MARS: Unleashing the Power of Variance Reduction for Training Large Models
D.6
Proof of Lemma C.5
Proof of Lemma C.5. By the definition of ne. it holds that
1 1
7 (s+ ( v 3(s + + 1)2/3 (s + n ne-
where the first inequality follows by the concavity of h2(x) = x//3, and the second inequality follows by 2 1, which implied that + t 2 and 27(s + - 1)2 1))²= 2 (s t)?. This finishes the proof.
E
Additional Experiment Results
E.1
Supplementary Results for the Main Experiments
Here we display the supplementary results for the experiments in Section The training and validation losses well as wall-clock time curves for small and medium models are displayed in Figures and And the O-shot and 5-shot evaluation results different downstream tasks for small, medium and large models listed in Tables 2,3,4 and Tables 5,6, respectively. It can be observed that different sizes of models trained with MARS-AdamY and MARS-Lion achieve better performances than baseline optimization methods with respect to cross-entropy loss, time efficiency well as downstream task performances.
Table 2. The evaluation results of small models pre-trained using the OpenWebTexx datasct (0-shot with wetevvlllaaion-ammmssss The best
scores each column are bolded. Abbreviations: HellaSwag HellaSwag, WG WinoGrande.
METHOD
ARC-E
ARC-C
BoolQ
HELLASWAG
OBQA
PIQA
WG
MMLU
ScIQ
AVG.
ADAMW
41.37
22.27
55.02
31.73
27.80
63.00
52.01
22.97
67.50
42.63
LION
40.15
21.93
59.72
31.72
26.00
62.95
51.07
22.92
64.80
42.36
MUON
39.73
23.55
57.31
30.84
25.00
61.48
50.36
22.89
62.70
41.54
MARS-ADAMW
40.70
23.63
59.17
32.46
27.00
61.92
51.22
22.98
67.40
42.94
MARS-LIO 40.78
23.72
51.74
31.59
29.20
62.68
51.30
22.94
65.50
42.16
Table The evaluation results of medium models pre-trained using the Webnee dataset (0-shot with Im- celluaatio-haaress, The best scores in each column are bolded. Abbreviations: HellaSwag = HellaSwag, WG
WinoGrande.
METHOD
ARC-E
ARC-C
BOOlQ
HELLASWAG
OBQA
PIQA
WG
MMLU
SCIQ
AVG.
ADAMW
43.43
23.98
58.13
37.76
27.20
65.56
52.49
22.80
67.60
44.33
LION
44.11
25.43
60.06
37.64
31.40
66.05
53.20
22.97
69.50
45.60
MUON
43.01
24.57
58.93
35.85
30.60
64.85
51.54
22.89
66.70
44.33
MARS- ADAMW
43.94
25.85
54.50
39.88
30.60
66.87
52.01
22.97
72.10
45.41
MARS-LION 45.33
24.74
55.84
38.80
30.60
64.96
53.83
23.33
68.70
45.13
Table The evaluation results of large models pre-trained using the WebnWeeText dataset (0-shot with haw.ovallaaaiioooaeesess
The best
scores each column bolded. Abbreviations: HellaSwag HellaSwag, WG WinoGrande.
METHOD
ARC-E
ARC-C
BoolQ
HELLASWAG
OBQA
PIQA
WG
MMLU
Sc1Q
AVG.
ADAMW
46.30
26.19
59.91
41.70
31.40
68.12
51.46
23.10
72.80
46.78
LION
47.73
26.45
$7.09
42.43
30.20
68.01
54.38
23.41
74.00
47.08
MUON
45.45
26.37
59.69
40.28
31.00
67.08
52.41
23.26
66.70
45.80
MARS-ADAMW
48.11
25.77
62.26
44.64
32.60
68.06
56.04
23.98
73.00
48.27
MARS-LION 47.77
26.71
59.45
43.07
31.20
68.39
55.72
24.53
72.50
47.70
Table 5. The evaluation results of small models pre-trained using the Open WebText dataset (5-shot with Im-evaluatioon-haress) scores each column are bolded. Abbreviations: HellaSwag HellaSwag, WG WinoGrande.
The best
METHOD
ARC-E
ARC-C
BOoLQ
HELLASWAG
OBQA
PIQA
WG
MMLU
SCIQ
AVG.
ADAMW
41.75
22.78
54.04
32.33
28.20
63.38
52.57
26.88
76.00
44.21
LION
42.21
22.70
55.41
31.82
24.80
62.40
53.04
24.63
74.80
43.53
MUON
41.50
23.46
48.78
30.48
24.60
61.26
52.01
24.63
67.20
41.55
MARS-ADAMW
45.24
24.66
56.97
32.76
25.60
62.40
50.43
25.78
76.70
44.51
MARS-LIC
43.06
22.78
55.66
32.17
26.20
62.24
50.59
25.32
72.80
43.42
E.2 MARS and MARS-approx.
We then conduct experiments to compare the performance of MARS and MARS-approx (MARS-AdamW instantiation) on GPT-2 small and medium models, the training and validation loss curves are shown in Figures and 5. Models trained with MARS exhibit consistently better performance than those trained with MARS- MARS-approx. This suggests that: (a) The exact version, which employs the variance reduction formulation, is more fundamental than the approximate version. (b) The
26

---

MARS: Unleashing the Power of Variance Reduction for Training Large Models
Table 6. The evaluation results of medium models pre-trained using the Weeeee dataset (5-shot with Im-evaluation--amees) The
best scores in each column are bolded. Abbreviations: HellaSwag = HellaSwag. WG = WinoGrande.
METHOD
ARC-E
ARC-C
BOOLQ
HELLASWAG
OBQA
PIQA
WG
MMLU
SCIQ
AVG.
ADAMW
48.23
25.43
45.26
38.32
27.60
65.83
52.33
26.21
80.90
45.57
LION
49.16
24.49
58.32
38.09
30.00
66.05
51.22
26.43
81.20
47.22
MUON
47.56
24.49
58.56
36.10
29.20
65.13
52.72
25.15
73.10
45.78
MARS- ADAMW
48.99
25.60
52.11
40.02
30.80
65.56
54.30
25.49
83,50
47.37
MARS-LION
49.03
25.77
$1.59
39.11
29.80
65.51
53.59
24.85
81.40
46.74
34
GPT-2 small (125M), openwebtext
3.4
GPT-2 small (125M), openwebtext
GPT-2 small (125M), openwebtext
Adammi
Adami
Lion
Lion
333
Pham
MAllS-AdamtW MARS Lion
MARlSAAaam MARSLion
MAPS Adumti
Jainn LOW
3.2
Valdazion
3.2
3.2
MARS-L in
3.1
SWllSettenn
3.3
3.0
3.0
3.0
2.9
2.8
20
30
40
55
2.8
10
30
60
50
2.8
5000
Thaining toilens
Thiining tobens (8)
clokk 36xA100 (seconds)
(a) Training Loss
(b) Validation Loss
(c) Wall-clock time
Figure 2. The training and validation loss curves, plotted against both training tokens and wall-clock time on GPT-2 small model (125M).
GPT-2 Medium (355M), opermebeext
GPT-2 Medium (355M), openwebtext
GPT-2 Medium (355M), openwebtext
32-
3.2
3.2
Adammi
Adami
Adamm
Utn
Lion
Lion
3.1
3.1-
MARS Adlamity
MARS AdaaW
MARS-AdimmV
3.0
MARS Lion
2.0
MARS Lion
3.0
MARtS-tion
29-
2.9
Validattoo
2.9
2.8
2.8
2.8
2.7
2.7
2.6
2.6
10000
20000
30000 40000 60000 32vA300 (seconos)
tobens
(a) Training Loss
(b) Validation Loss
(c) Wall-clock time
Figure 3. The training and validation loss curves, plotted against both training tokens and wall-clock time on GPT-2 medium model (355M).
approximate version serves as a practical alternative in scenarios where computational efficiency is a priority, as it incurs only minimal performance loss. However, in settings where maximizing validation accuracy is crucial, the exact version is recommended.
E.3 Experiments on Wed-Eue 100B Dataset
FineWeb-Edu dataset (Lozhkov al., 2024) is a high-quality dataset based on well-filterdd educational web pages. To better investigate the efficiency of our algorithm, we also use Fine Web-Euu 100B, a subset of FineWeb-Edu with around 100B tokens to train GPT-2 small (125M) and XL (1.5B, with the same learning rates as GPT-2 large models) models with optimizers including Adam W. Muon and MARS- MARS-AdamW-approt W-approx. We leave around B. tokens for validation and other tokens for training. The training and evaluation curves are shown in Figures and 7. It can be seen that our algorithms can also achieve better performances even with different datasets. For a comprehensive investigation, we evaluate these models on metrics same as experiments in Section 4.2, and the results are shown Tables 7 and 8. We also compare the results with the open-source GPT-2 models on Hugging Face (Radford al., 2019) (denoted as "OpenAI-Commi" in the tables). Compared with Table 2, it can be observed that this dataset is actually better for the superior performances. However, models trained with our algorithm can also show advantages over baseline optimization approaches trained with such a high-quality dataset.
27

---

MARS: Unleashing the Power of Variance Reduction for Training Large Models
GPT-2 small (125M), openwebtext
GPT-2 medium (355M), openwebtext
2.96
MARS-approx MARS-exact
MARS-approx
MARS-exact
2.94
2.92
2.68
Training Loss
Training Loss
2.90
2.66
2.88
2.64
2.86
2.62
2.84
2.60
2.82
2.58
40
10
30
50
Training takens (B)
Training tokens (B)
Figure 4. Training loss curves for W and MARS-Adam Aaas--appoap GPT-2 small (125M, left) and medium (355M, right). pretrained with Open Wex dataset and plotted against training tokens.
GPT-2 small (125M), openwebtext
GPT-2 medium (355M), openwebtext
3.00
MARS-approx
MARS-approx
2.98
MARS-exact
MARS-exact
2.74
2.96
Validation Loss
2.94
Validation Loss
2.72
2.70
2.68
2.90
2.88
2.66
2.86
2.64
20 Training tokens (B)
40
10
30
Training tokens (B)
Figure 5. Validation loss curves for MARS- and MARS-AdamW-aprroxx GPT-2 small (125M, left) and medium (355M, right),
pretrained with OpenWebTecd dataset.
Table The evaluation results of small models pre-trained using the FineWeb- Edu 100B dataset (0-shot with Im nn-evaluation evaluation-harness) The
best scores in each column are bolded. Abbreviations: HellaSwag HellaSwag, WG WinoGrande.
METHOD
ARC-E
ARC-C
BoolQ
HELL-ASWAG
OBQA
PIQA
WG
MMLUU
SCIQ
AvG.
OPENAL- PPENAL-COMM.
39.48
22.70
48.72
31.14
27.20
62.51
51.62
22.92
64.40
41.19
ADAMW
51.43
26.54
55.78
36.26
30.60
64.53
50.36
24.49
71.50
45.72
MUON
47.85
27.56
57.16
33.46
31.60
63.66
51.30
23.17
67.30
44.78
MARS-ADAN ADAMM
52.23
27.39
55.84
36.91
32.20
64.80
49.96
22.95
71.10
45.93
Table evaluation results models pre-trained using the FineWeb-Edu 100B dataset (0-shot with Im-evaluatioon-hames) The best scores in each column are bolded. Abbreviations: HellaSwag HollaSwag, WG WinoGrande.
METHOD
ARC-E
ARC-C
BOoLQ
HELLASWAG
OBQA
PIQA
WG
MMLU
SCIQ
AvG.
OPENAT-COMEL
51.05
28.50
61.77
50.89
32.00
70.51
58.33
25.24
76.00
50.48
ADAMW
68.22
38.40
61.13
53.93
39.00
72.69
54.78
25.47
85.30
55.43
MUON
64.18
36.52
58.38
51.83
37.40
72.03
55.56
24.93
81.90
53.64
MARM-
66.54
39.85
63.82
56.52
41.20
73.34
$6.59
23.86
86.00
56.41
E.4 Computer Vision Experiments
We also carry out experiments on the classification task in the field of computer vision. We conduct experiments with ResNet-18 18 model (He et al., 2016) on the CIFAR-I0 and CIFAR-100 datasets (Krizhevsky et al., 2009) for Adam Lion,
28

---

MARS: Unleashing the Power of Variance Reduction for Training Large Models
GPT-2 small (125M), FineWeb-edul008
GPT-2 XL (1.5B), FineWeb-edu1008
2.9
AdamW
AdamW
Muon
Muon
MARS-AdamW
2.8
MARS-AdamW
Training Loss
3.2
Training Loss
2.7
2.6
3.1
2.5
2.3
10
20
30
40
50
10
20
30
40
50
Training tokens (B)
Training takens (B)
Figure 6. Training loss curves for AdamW, Muon and MARS-A AaRS-AAdm--appa GPT-2 small (125M, left) and XL (1.5B, right), pretrained with Fine Web-edu 100B dataset and plotted against training tokens.
GPT-2 small (125M), FineWeb-edul008
GPT-2 XL (1.5B), FineWeb-edu1008
2.9
AdamW
AdamW
Muon
Muon
3.3
MARS-AdamW
2.8
MARS-AdamW
Validation Loss
3.2
Validation Loss
2.7
3.1
2.6
2.5
2.4
2.8
2.3
10
30
40
10
20
50
Training tokens (B)
Training tokens (B)
Figure 7. Validation loss for AdamW, Muon and MARS-AdamW-appoor GPT-2 small (125M, left) and XI. (1.5B, right). pretrained with FineWeb-edu 100B dataset and plotted against training tokens.
Shampoo' Muon, and variants of MARS instantiation, following the setting in Chen al. (2018).
We do grid search explore the best hyper-parameters for each of these optimization methods. We search over {10-5 ..., 10º) for the learning rate and {0, ...: 1.0} for the weight decay. We set 31 0.9 and search over {0.99, 0.999) for 33 for AdamW and Muon; fix 31 0.9 and search Bg over f0.99, 0.999) for Lion; search over {0.9, 0.95) for 8n and 0.025 for MARS models. We train
(0.95, 0.99, 0.999) for 32 for Shampoo; and we fix the (51, }2) = (0.95, 0.99) and for 200 epochs with training batch size 128 on NVIDIA A6000 GPU. And we also apply MultiStepLR scheduler so that the learning rate would decrease to 10% of the original rate at the 100th epoch and to 1% at the 150th epoch. We display the test loss and test accuracy for CIFAR-I0 10 and CIFAR- 100 datasets Figures and 9, respectively. The results show that our algorithm can achieve better validation loss after the decay of learning rate and better test accuracy within the final stage of training.
We also compare the performances among the baselines of Adam W. Lion and Shampoo without variance reduction, the approximate and the exact versions of MARS instantiations. The test loss and accuracy curves for CIFAR 10 dataset are shown in Figures, and Figures, respectively. And the test loss and accuracy curves for CIFAR100 dataset are shown in Figures 10-11, 11, and Figures 12-13, respectively. Moreover, observed that the exact versions perform little better than the approximate versions, but much better than the baseline approaches, showing the superiority of variance reduction in MARS.
list the best test losses and accuracies in Table 9. It can be
In practice, Distributed Shampoo (Shi et al., 2023) to facilitate the training of Shampoo.
29

---

MARS: Unleashing the Power of Variance Reduction for Training Large Models
CIFAR10
CIFAR10
AdamW
96
0.9
Lion
Shampoo
94
0.8
Muon
Test Accuracy (%)
92
Test loss
0.7
MARS-Lion
90
MARS-Shampoo
0.6
88
AdamW
0.5
86
Shampoo
0.4
Muon
84
MARS-AdamW
82
MARS-Lion
MARS-Shampoo
0.2
25
125
150
175
50
75
125
150
175
200
Epochs
Epochs
(a) Test Loss
(b) Test Accuracy
Figure 8. The test loss and test accuracy for different optimizers CIFAR-10 10 dataset.
CIFAR100
CIFAR100
4.0
AdamW
Lian
3.5
Shampoo
Muon
70.
MARS-AdamW
Test Accuracy (%)
MARS-Lion
Test loss
2.5
MhapoSh
65
AdamW
Lion
60
Shampoo
Muon
1.5
MARS-AdamW
MARS-Lion
1.0
MARS-Shampoo
50
25
50
75
100 Epochs
125
150
175
200
50
75
100
125
150
175
200
Epochs
(a) Test Loss
(b) Test Accuracy
Figure The test loss and test accuracy for different optimizers on CIFAR- dataset.
Table The best test losses and accuracies of different optimizers on CIFARIO and CIFARI00 datasets during the training epochs. The bolded are the best test loss or best test accuracy among the listed optimizers.
Best Test Loss
Best Test Accuracy
Optimizer
CIFARIO
CIFARI00
CIFARIO
CIFAR100
AdamW
0.230
1.726
94.81
73.70
Lion
0.245
1.351
94.68
74.28
Shampoo
0.354
2.426
94.65
74.27
Muon
0.306
2.608
95.08
74.64
MARS-A AdamW-apppn
0.199
0.971
95.29
76.97
MARS-ACH AdamW
0.193
0.888
95.26
77.38
MARS MARR--Lin-aapp aoon-approx
0.202
0.985
95.05
75.97
MARS-Lion
0.219
0.991
94.98
76.15
MARS- MARS-Shampoo--ppprx
0.202
1.256
94.92
74.80
MARS-
0.194
0.982
94.98
75.83
E.5 Sensitivity to 7
To explore the impact of 7t, we test various ys on GPT-2 small model, including constant and linearly changing schedules. And we plot the training and validation curves in Figure 14. can be observed that there are slight differences among
30

---

MARS: Unleashing the Power of Variance Reduction for Training Large Models
1.0
1.0
1.0
Adamm
Uon
Shampoo
0.9
0.9
MARSL iem-pprss MARS4.Kn kon
0.9
M6t5-Shempo--pppm
MARtSkddemi
KARS-Shempio
0.3
6.8
0.8
0.7
6.7-
0.7
0.6
0.6
0.6
0.5
0.5-
0.5
0.4
6.4
0.4
0.3
0.3
0.3
0.2
6.2
0.2
25
50
100 Epocts
125
150
175
25
50
75
100 Epochs
125
150
175
200
25
50
75
100 Eprocts
1.25
1.50
175
200
(a) Adam W-seiie
(b) Lion-series
(c) Shampoo-series
Figure 10. The test loss curves for the baselines of AdamW, W. Lion and Shampoo without variance reduction, the approximate and the exact versions of MARS instantiations CIFAR-I0 dataset.
GIAR10
CIFARIO
OFAR1
90-
56
Lion
96-
MARS-Loon-eprox
94
94
MARS
04
92-
92-
92-
60-
Acqurccy
50-
90-
88
Adamlt
Shampoo
92-
MARS AdaniM aporox
82-
82-
MARS Shampoo аporоx
MARS AdamiN
MARS Shampco
80
80
150
100
100
Epochs
Epochs
Epochs
(a) AdamA-dassss W-series
(b) Lion-series
(c) Shampoo-series
Figure 11. The test accuracy curves for the baselines of AdamW, Lion and Shampoo without variance reduction, the approximate and the exact versions of MARS instantiations on CIFAR-IO 10 dataset.
CIFARID0
CIFAR100
CIFARIDO
Adamwi
Lion MARSLion
Shampoo
MAR5-AdamW
аррrок
3.5
MARS-Shampos
Shampoo r0
3.0
3.0
3.0
1044
2.5
2.5
2.5
2.0-
2.0
2.0-
1.5
1.5
1.5
1.0
1.0
25
50
100 Epochs
125
150
175
200
25
50
75
100 Epochs
125
150
175
200
0
25
50
75
100 Epochs
125
150
175
200
(a) Adam W-series
(b) Lion-series
(c) Shampoo-series
Figure 12. The test loss for the baselines of AdamW, Lion and Shampoo without variance reduction, the approximate and the versions of MARS instantiations on CIFAR- 100 dataset.
different ys where 0.025 is the best 7. Therefore, we used 7 = 0.025 for other experiments in this paper.
E.6 Different Learning Rate Scheduler
E.6.1
CONSTANT LR
To ensure a fair comparison by eliminating the influence of learning rate changes during training and explore the potential for continuous training with MARS, conduct supplementary experiments GPT-2 small, medium, and large models using a constant learning rate for both AdamW and MARS-Adam waaproppoox. For each group of experiments, we compared between 2 different maximum learning rates. The training and validation curves are displayed in Figures 15 and 16. The results also indicate that our algorithm has superior performances over Adam W optimizer under a fair comparison of
31

---

MARS: Unleashing the Power of Variance Reduction for Training Large Models
CIFAR100
Lion MARS Lion approx
MARS
Accuncy (5)
Accuracy
Accurecy
65-
60-
55-
Shamoco
NARS-AcamiiN-pppm KARS-AdamtN
MARS- SMRRS-Shampo--ppee aporos
M0RR--Shampon
50
150
125
150
175
50
50
75
100 Epocho
12s
150
17s
200
50
150 Epochs
125
150
175
(a) Adam W-series
(b) Lion-series
(c)
Shampoo-series
Figure 13. The test accuracy curves for the baselines of AdamW, Lion and Shampoo without variance reduction, the approximate and the exact versions of MARS instantiations on CIFAR-100 dataset.
3.000
GPT-2 small (125M), openwebtext
3.00
GPT-2 small (125M), openwebtext
y= 0.0001
Y- 0.0001
2.975
2.98
y= 0.000
0.0.01
y 0.01
2.950
y- 0.025
2.96
y= 0.025
y- 0.05
Validation Loss
y= 0.05
Training Loss
2.925
==
2.94
0.1
y 0.2
0.2
2.900
2.875
2.88
2.86
2,825
10
20
30
40
10
20
30
40
50
Training (B)
Training
(B)
Figure 14. Training loss and validation loss curves for MARS and MARS-approp on GPT-2 small (125M) models trained with MARS- AdamW-approx with different 78, pretrained with Open Web dataset and plotted against training tokens.
constant learning rates.
GPT-2 openwebtext
(355M), openwebbext
Large openwebtext
3.5
3.3
3.2
AdantN-LR- R-Bed
Adem
AdwmWAAR-
MAlts-1t-bb-4
MAlIS-B--Jeed
AdamWVALit land
Adanti-LiR-33
AdamW-Litw L.5
3.1
MARS 2it2204
-AR---8-
MARS-LR-1. 5m-3
MARS--R-iiw-
3.3
Trahing Look
3.1
3.0-
3.0
Tratiing
2.9
2.8
3.9
2.7
30
40
50
10
20
30
40
So
10
20
Training tollers (B)
Thaining tokens (B)
Thaining tokkens
Figure 15. Training loss curves for MARS and MARS-approx on GPT-2 small (125M, left), medium (355M, middle) and large (770M, right) with constant learning rate, pretrained with Weenee dataset and plotted against training tokens.
E.6.2 WSD SCHEDULER
Cosine learning rate scheduler is utilized in most of experiments. Recently Hu et al. (2024) introduced novel learning rate scheduler called
Siouup-Stabe--eeece (WSD) scheduler, which composed of 3 stages, including learning rate warmuu stage, constant learning rate stage, as well as learning rate decay stage. To test the flexibility and ability of continuous training of MARS trained with respect to learning rate schedulers, we implement experiments on GPT-2 small and medium models with Adam W and MARS-Adam W-approx scheduled with WSD for 10k, 20k, 50k and 100k steps. The
32

---

MARS: Unleashing the Power of Variance Reduction for Training Large Models
GPT-2 small (125M), openwebtext
GPT-2 Medium (355M), opernwebtext
GPT-2 Large (770M), openwebtext
Saami-LR=3e-
AeamWAR-224
3.4
3.2-
Adamw/ LR=1.533
31-
M6RS-UR=3e-
MARS 5e
MARS LR-le3
2.0
Valicttonn
Valtiatem
32-
3.0
29
2.8
2.8
2.7
2.9
Trining
Trinning
Figure 16. Validation loss for MARS and MARS-approx GPT-2 small (125M, left), medium (35SM, middle) and large (770M, right) with constatnt learning rate, pretrained with OpenWeb'Tar datasct and plotted against training tokens.
training and validation loss curves shown in Figures 17 and 18. The curves indicate that MARS has better potential for continuous training and exhibits an explicit edge over baseline algorithm. GPT-2 small (125M), openwebtext
GPT-2 Medium (355M), openwebtext
3.4
3.2
AdamW-LR=Gon
AdamW-LR-3e-
AdamW-LR=-20
AdamW-LR- 1=-15e
3.3
MARS-AdamW
3.1
MARS-Adamin
3.2
3.0
Training Loss
Training Loss
2.9
3.1
2.8
2.7
2.6
2.8
10
20
30
40
50
10
20 Training takens (B)
40
50
Training tokens (B)
Figure 17. Training loss curves for AdamW (with different maximum learning rates, labeled with "LR") and AAARS-AdamWar GPT-2 small (125M, left) and medium (355M, right) with WSD Scheduler for 10k, 20k, 50k and 100k steps, pretrained with Open WebTeet dataset and plotted against training tokens. GPT-2 small (125M), openwebtext
GPT-2 Medium (355M), openwebtext
3.2
AdamW-LR=6e-4
AdamW-LR=3e-4
AdamW-LR=3e-3
3
MARS-AdamW
MdRM-AdamW
Validation Loss
3.2
Validation Loss
3.0
3.1
2.9
3.0
2.8
2.9
2.7
10
Training tokens (B)
Training takens
Figure Validation curves for AdamW (with different maximum learning rates, labeled with "LR") and MARS- MARS-AdamW AdamW-approx on GPT-2 small (125M, left) and medium (355M, right) with WSD Scheduler for 10k, 20k, 50k and 100k steps, pretrained with OpenWehTes. dataset and plotted against training tokens.
33

---

MARS: Unleashing the Power of Variance Reduction for Training Large Models
E.7 Sensitivity to Batch Size
We also investigate the sensitivity batch size of algorithm. We implement experiments with MARS-AdamW GPT-2 small with batch size 240, 480 960, and compare them to AdamW. We fix other hyper-parameters the same the main experiments and trained with 80,000 steps. The results are plotted in Figure 19. It can be seen that MARS-Adaml consistendly outperforms AdamW, with the performance gap widening at smaller batch sizes supporting the intuition that variance reduction is especially effective in high- (small-batch) settings.
GPT-2 small (125M), openwebtext
GPT-2 small (125M), openwebtext
3.5
3.5
AdamW-240
AdamW-240
MARS-AdamW
3.4
MARS-AdamW-244)
AdamW-480
AdamW-480
MARS-AdamW
MARS-AdamW-480
3.3
AdamW-966
AdamW. 960
Loss
3.2
MARS Adam-AdaaWaaa 9-0
3.2
MARS-AdamW-onn
Validation
Train
3.1
3.1
3.0
2,9
2.9
2.8
2.8
10
30
40
50
60
70
80
10
20
30
50
60
70
Iterations (x10" steps)
Iterations (x10² steps)
Figure 19. The training and validation loss curves for varying global batch sizes (240/480/960), plotted against iteration steps on GPT-2 small model (125M).
Hlyper-parameter Settings
For training parameters, we did a grid search over learning rates between - 4, -1e-4,
1.5c 4, 4, 6e 4,1e - 3, 1.5e
3, 3e 3, 6e -3.. for weight decay coefficient, we did a grid search over {1e 1, 1e 2, le -33- For AdamW baseline, although we utilized the golden standard learning rates in literature (a parameter search for Adam have been done in Liu al. (2023)), we also did a grid search on different parameters. Part of the results of different learning rates and learning rate schedules are also shown in Section E. and E.6.2, respectively. Table 10 summarizes the architectural hyperparameters for GPT-2 models with 125M (small), 355M (medium), 770M (large) and 1.5B (XI., only used in Appendix E.3) parameters. Table 11 lists the general hyperparameters used across all experiments, while Table 12 present the training hyperparameters for the small, medium, and large models, respectively.
Table 10. Architecture hyperparameterss for GPT-2 series models (Radford al., 2019).
Model
#Param
#Layer
Tibead
demb
GPT-2 small
125M
12
12
768
GPT-2 medium
355M
24
16
1024
GPT-2 large
770M
36
20
1280
GPT-2 XL
1.5B
48
25
1600
Table 11. General
for the experiments.
Hyper-parameter
Value
Steps
100,000
Batch size in total
480
Context length
1024
Gradient clipping threshold
1.0
Dropout
0.0
Learning rate schedule
Cosine
Warm-up steps
2000
Base seed
5000
34

---

MARS: Unleashing the Power of Variance Reduction for Training Large Models
Table 12. Hlyper-parameters for GPT-2 experiments. use "/ 0.025 for MARS and μ 0.95 for Muon.
Hyper-parameter
GPT-2 Size
Adam
Muon
MARS-AdamW
small (125M)
6e-4
2e-2
6e-3
Max learning rate
medium (355M)
3e-4
le-2
3e-3
large (770M)
2e-4
6.67e-3
2e-3
small (125M)
3c-5
3e-5
3e-5
Min learning rate
medium (355M)
6c-5
6c-5
6c-5
large (770M)
lc-5
le-5
le-5
(0 32)
small//mciium/lage
(0.9,0.95)
(0.95,0.99)
35


Reviews:
[{"review_id": "LRuZxLNXgl", "reviewer": "Reviewer_8Y9B", "summary": "This paper proposes MARS, which is a framework that applies variance reduction to preconditioning methods for stochastic optimization. Namely, MARS takes a scaled update from STORM (which calculates momentum as an EMA combining the stochastic gradient and a gradient correction term) and applies preconditioning on the resulting variance-reduced object. The preconditioning can be achieved with different methods performing different approximations of the Hessian, leading to different instances of MARS-optimizers. The authors show that MARS has a superior convergence rate to AdamW. They empirically show that MARS-AdamW(-Approx) and MARS-Lion(-Approx) show competitive performance to Muon, Lion, and AdamW on GPT-2 trained on OpenWebText/FineWeb-Edu and ResNet-18 on CIFAR-10/100.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thorough review and constructive feedback. We sincerely appreciate your recognition of the novelty of our proposed framework and the clarity of our writing. Additionally, your detailed questions have been invaluable in helping us refine our work, and we listed our responses to your questions below.\n\n---\n**Q1**: I'm not sure if 'unified' is the right word here, because variance reduction is orthogonal to preconditioning\n\n**A1**: Thank you for your comment. We agree that variance reduction is orthogonal to preconditioning, and we appreciate the opportunity to clarify our use of the term. Our use of “unified” refers to two key aspects of the MARS framework: (1) The scaled variance-reduced momentum we propose not only generalizes STORM-style updates, but also recovers and accommodates various forms of momentum used in practice, including Nesterov, heavy-ball, Adan, and Lion. This highlights a shared underlying structure that our analysis brings to light. (2) The framework is designed to be applied on top of a wide range of preconditioning strategies, including those used in first-order methods like AdamW as well as second-order approaches like Shampoo. While variance reduction and preconditioning are orthogonal, MARS's integration enables theoretical analysis across both settings.\n\nTo the best of our knowledge, this is the first framework that simultaneously accommodates optimizers as diverse as STORM and Muon under a common formulation. We will revise the relevant paragraphs in the manuscript to better reflect this intent.\n\n---\n**Q2**: Are there characteristics of settings in which variance reduction should help\n\n**A2**: Thank you for your insightful question. It is well known that deep learning training—particularly for large language models—exhibits high gradient variance (e.g., McCandlish et al., 2018; RAdam, Liu et al., 2020). MARS introduces key design changes—such as a scaled, tunable correction mechanism—that make it more effective in this setting.\n\nWe believe it is these design choices, rather than the setting itself, that drive the improved performance. For a detailed discussion of the algorithmic differences, please refer to our response to Reviewer ZHwA (A3).\n\n---\n**Q3**: The paper would benefit by bringing some of the results to the main paper and deferring details about the instantiations\n\n**A3**: Thank you for your suggestions. Due to space constraints, we were limited in the amount of experimental results we could include in the main paper. In the final version, we will restructure the content to highlight more key results in the main paper.\n\n---\n**Q4**: Typos\n\n**A4**: Thank you for spotting these typos. We have corrected them and will ensure they are reflected in the final version.\n\n---\n**Q5**: What was the reasoning behind the choice in not including MARS-Shampoo for the language model pretraining experiments?\n\n**A5**: The time efficiency of Shampoo is too low due to the SVD decomposition. For CIFAR100 experiments, Shampoo requires 14.75 hours while AdamW only consumes 41 minutes for 200 epochs. For GPT-2 small models, we noticed that training with Shampoo is more than 5 times slower than AdamW experiments. While we acknowledge that accelerated implementations of Shampoo could mitigate this issue, we were unable to find a reliable open-source version to use in our experiments.\n\n---\n**Q6**: Learning rate stability of MARS-AdamW\n\n**A6**: Thank you for the question. As shown in our learning rate ablation study (Figure 16), MARS-AdamW demonstrates greater stability compared to AdamW across a range of constant learning rates. Notably, its advantage becomes more pronounced at higher learning rates. This suggests that MARS-AdamW can tolerate and benefit from a larger maximum learning rate than AdamW.\n\n---\n**Q7**: Sweep across batch size\n\n**A7**: We conducted experiments across multiple batch sizes (240, 480, 960) as suggested. Using the difference between AdamW and MARS-AdamW as a measure of improvement, we included the loss curve of varying batch sizes in in Figure 1 in https://anonymous.4open.science/r/M-E6D7/F.pdf. As shown below, MARS-AdamW consistently outperforms AdamW, with the performance gap widening at smaller batch sizes—supporting the intuition that variance reduction is especially effective in high-variance (small-batch) settings. This conclusion also aligns with our findings in Figure 16 (initial submission) under large step sizes.\n\n| AdamW - MARS-AdamW | Train loss | Valid Loss |\n| ------------------ | ---------- | ---------- |\n| 240                | 0.0477     | 0.0359     |\n| 480                | 0.0406     | 0.0327     |\n| 960                | 0.0384     | 0.0322     |\n\n---\n**Q8**: Compare to methods like AdEMAMix\n\n**A8**: : We included tuned AdEMAMix as an additional baseline in Figure 3 in https://anonymous.4open.science/r/M-E6D7/F.pdf. While it benefits from longer-term averaging, it performs worse than MARS in both train and validation loss.", "claims_and_evidence": "It is clear that MARS can be applied on any choice of adaptive gradient method, and the authors show that adding MARS on AdamW and Lion results in improved performance on language model pretraining at various scales (125m, 355m, 770m), language model finetuning (eg. Table 1) and computer vision tasks (Appendix E.4).\n\nThis may be pedantic, but the authors say numerous times that MARS is a framework which ‘unifies’ first-order and second-order optimization methods (eg. Lines 114-116, column 2). I’m not sure if ‘unified’ is the right word here– it seems that MARS is simply applying a scaled version of the STORM momentum, which you can then perform any preconditioning method on top. I wouldn’t say that it ‘unifies’ the first and second order methods, but is rather an addition that can be applied on top of the adaptive gradient method of your choice, because variance reduction is orthogonal to preconditioning.", "methods_and_evaluation_criteria": "The authors cover a variety of datasets tracking reasonable metrics for pretraining and computer vision tasks. For pretraining, they also show that downstream performance is also improved compared to other optimizers when controlling for number of pretraining steps. Implementing MARS in practice typically needs two forward passes on two different batches, but the authors conduct sweeps over hyperparameters (particularly learning rate, $\\gamma$) and their -exact and -approx version (which replaces $\\nabla f(x_{t-1}, \\xi_t)$ with $\\nabla f(x_{t-1}, \\xi_{t-1})$ in the MARS update) to show that performance is comparable.", "theoretical_claims": "I have not rigorously verified the convergence guarantees. It seems that it is a relatively straightforward extension of the STORM approach.", "experimental_designs_or_analyses": "Experimental design and analyses are sound. I believe some additional analyses would strengthen the paper, given that the proposed framework combines existing variance reduction methods with existing adaptive gradient methods. More specifically, it was previously suggested that variance reduction can be ineffective in deep learning but the authors find that it can be beneficial here; are there characteristics of settings in which variance reduction should help, or are there characteristics of the adaptive gradient method that may yield greater empirical benefit with MARS as opposed to other adaptive gradient methods? Furthermore, the authors (please correct me if I'm wrong) use a fixed batch size across model scales, but an analysis between small-batch vs. large-batch training would be insightful, given that variance reduction intuitively should help more in smaller-batch or high-variance scenarios. Finally, the authors report multiple baselines (AdamW, Lion, Muon) but do not have other baselines which yield similar variance reduction effects (see Questions).", "relation_to_broader_scientific_literature": "There is an established line of variance-reduced methods as mentioned by the authors (eg SVRG, SARAH, SPIDER, STORM) as well as adaptive gradient methods. The authors show that variance-reduction can indeed help even at scale—they cite prior negative results on variance reduction for deep neural nets (e.g., Defazio & Bottou’s arguments) and give evidence that in certain modern settings, variance reduction can indeed be beneficial.", "essential_references_not_discussed": "None; see Questions for AdEMAMix citation.", "other_strengths_and_weaknesses": "* The paper is generally well-written and presents prior background on variance reduction techniques well.\n* The paper demonstrates that variance reduction techniques can be added as a separate dimension to preconditioning for stochastic optimization and yield further benefits in practical settings.\n\n* I believe a lot of space is taken in the main paper simply reiterating the updates of known optimizers (AdamW, Lion, Shampoo) and replacing the momentum term with the scaled gradient correction term; there are very few results in the main paper, and many in the appendix. I would suggest to the authors that the paper would benefit bringing some of these results (such as ablating -approx vs -exact, one vision experiment, etc.) to the main paper and deferring details about the instantiation of MARS to the appendix.\n* For novelty, the paper combines two known techniques (STORM modulo a scalar correction, and adaptive optimization methods), although I see the value in the empirical experiments showing that variance reduction techniques have potential in further improving language model performance which is compatible with preconditioning. Thus, I lean towards acceptance of the paper.", "other_comments_or_suggestions": "* Line 092, first column: “variance-reduction” -> “variance reduction”\n* Line 163, first column: “precondtioned” -> “preconditioned”\n* Line 1476: \"grid research\" -> \"grid search\"\n* Figure 13: y-axis should be \"Test Accuracy\"?", "questions_for_authors": "1. What was the reasoning behind the choice in not including MARS-Shampoo for the language model pretraining (even just GPT-2 small) experiments?\n2. Are there any insights regarding the learning rate stability of MARS-AdamW? If the authors could provide some insight about tuning the MARS-AdamW learning rate (eg. it seems that it is typically higher than the optimal LR for the adaptive method chosen) this would be a valuable addition to the paper.\n3. Did the authors sweep across batch size for the pretraining experiments? Does MARS show less pronounced improvements in the large -batch setting (even on GPT-2 small?)\n4. How do the variance reduction techniques employed by MARS, STORM, etc. compare to methods like AdEMAMix[1] which simply maintains multiple EMAs with longer retainment of the past?\n\n[1] Pagliardini, Matteo, Pierre Ablin, and David Grangier. \"The ademamix optimizer: Better, faster, older.\" arXiv preprint arXiv:2409.03137 (2024).", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "s2xWC4GNPR", "reviewer": "Reviewer_L7rR", "summary": "This paper studies variance-reduction algorithm for large language models, and proposes a general framework, MARS, that combines variance-reduction and pre-conditioning. In particular, the authors add a control parameter $\\gamma_t$ to the STORM momentum aggregate, i.e., $m_t = \\beta m_{t-1} + (1-\\beta) c_t$ where $c_t = \\nabla f(x_t,\\xi_t) + \\gamma_t \\frac{\\beta}{1-\\beta}(\\nabla f(x_t,\\xi_t) - \\nabla f(x_{t-1},\\xi_{t-1}))$. $\\gamma_t$ provides interpolation between noise control and bias: setting $\\gamma_t=1$ recovers the variance-reduction update of STORM, and setting $\\gamma_t=0$ recovers standard SGDM. In addition, the authors also incorporate the pre-conditioning technique to the variance-reduced momentum, namely $x_{t+1} \\approx x_t - \\eta_t H_t^{-1}m_t$, where $H_t$ is some approximation (either full-matrix or diagonal) of the hessian $\\nabla^2 F(x_t)$. Upon substituting popular pre-conditioning methods in the literature, the authors instantiate different algorithms of MARS. The authors provide theoretical analysis of the convergence guarantee of the proposed framework, and presented empirical evidence that MARS have better performance than the state-of-art optimizers on training language models.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed and thoughtful review. Below, we provide detailed responses to each of your comments.\n\n---\n**Q1**:The name MARS-Shampoo\n\n**A1**: The interpretation of Shampoo as a projection onto the closest semi-orthogonal matrix ($UV^\\top$) originates from Proposition 5 of [1], which presents Shampoo as steepest descent under the spectral norm. While this interpretation is not directly proposed in Gupta et al., 2018, our MARS-Shampoo algorithm builds on (19) from [1]. Moreover, MARS-Shampoo employs Newton iteration, in contrast to the Newton-Schulz method used in Muon. We will clarify that this is an approximate version of Shampoo in our revision.\n\n[1] Bernstein et al., Old optimizer, new norm: An anthology\n\n---\n**Q2**:The theorems prove convergence of MARS without clipping\n\n**A2**: You are correct. This discrepancy arises because our theoretical results rely on standard smoothness assumptions, which simplify the analysis but do not directly account for the clipped version.\n\nIn practice, gradient clipping is employed as a heuristic to enhance stability, particularly in deep learning where gradients may explode. While our work does not theoretically justify the clipped variant, prior studies (e.g., [1]) have shown that clipping can handle scenarios where gradient smoothness conditions are violated or grow with gradient norms. We leave the extension of our analysis to incorporate clipping—and its interplay with momentum and adaptive gradient methods—as future work.\n\n[1] Zhang et al., Why gradient clipping accelerates training: A theoretical justification for adaptivity.\n\n---\n**Q3**:Typos\n\n**A3**: Thank you, we have corrected them.\n\n---\n**Q4**:Experimental difference between Muon's report and Figure 1 of our paper\n\n**A4**: The performance difference primarily stems from differences in experimental setups: while Muon’s speedrun combines architectural changes with its optimizer, our evaluation isolates the optimizer’s impact using Muon's open-sourced implementation on GPT-2 architecture. We rigorously tuned learning rates, but differences persist due to factors like dataset choice (FineWeb vs. OpenWebText), training scale (<10B vs.50B tokens), and learning rate schedulers (WSD vs. Cosine). Regarding the MARS-Shampoo algorithm, we added experiments on CV tasks with varying γ (see Table 1 in https://anonymous.4open.science/r/M-E6D7/F.pdf). For CIFAR-100, a larger value γ>1-μ (0.1) yields the best performance, while for CIFAR-10, a smaller value γ<1-μ (0.001) performs best.\n\n---\n**Q5**:Is Lemma D.1 a new finding.\n\n**A5**: Thank you for your interest in Lemma D.1. While we independently derived this connection, we later found that similar ideas have appeared implicitly in prior work. A special case of the forms in (D.1) and (D.2) align with the implementation of Nesterov in PyTorch, and a special case of the 2nd form is stated as AGD-II in Lemma 1 of Adan [Xie et al. 2022]. However, to the best of our knowledge, Lemma D.1 is the first to present a unified formulation and theoretical framework that encompasses both classical momentum and noise-corrected gradients—capturing HB, NAG, Lion, and STORM within a single view.\n\n---\n**Q6**: How do Theorem B.4 and B.5 imply a convergence rate\n\n**A6**: As a sketch of proof, after Jensen's inequality, we sum up $||x_{t+1} - x_t||/\\eta_t = ||H_t^{-1} m_t|| $ with $||\\nabla F(x_t)-m_t||/\\rho \\ge \\frac{1}{||H_t||} ||\\nabla F(x_t)-m_t||$. Triangle inequality (on $m_t$ and $\\nabla F(x_t) - m_t$) and summing over $t$ leads to a bound of the form $\\frac{1}{||H_t||} ||\\nabla F(x_t)||\\le O(T^{−1/3})$. While the final convergence requires upper boundedness of $||H_t||$, this is commonly satisfied in practice. We will include the detailed derivation in the final revision of our manuscript.\n\n---\n**Q7**: Empirical performance of Nesterov AdamW (or Adan). Also, could the authors further elaborate on γ≈1-μ?\n\n**A7**: Thank you for the feedback. We compared MARS with Adan and found that while Adan achieves similar final performance under a cosine schedule, it consistently underperforms MARS at the same learning rate (see Figures 5–6 in the link above).\n\nFor the choice of γ, we note that the performance of DL tasks can be influenced not only by theoretical factors but also by practical considerations such as smoothness and stability. These factors can hinder performance as γ increases. To validate the role of variance reduction, we examine it from two perspectives:\n\n(a) Comparing exact vs. approximate variance reduction (Figure 5). and\n\n(b) We explore γ values not strictly tied to 1-μ (e.g., γ=0.025 vs. μ=0.95).\n\nWe added experiments for MARS-AdamW on CV tasks with varying γ (see Table 2 in the link above). For CIFAR-10, a larger value γ>1-μ (0.1) yields the best performance, while for CIFAR-100, a smaller value γ<1-μ (0.01) performs best.\n\n---\n**Q8**: MARS-SGDM\n\n**A8**: We have now included a comparison between MARS-AdamW and MARS-SGDM in Figure 4 in the link above.", "claims_and_evidence": "yes", "methods_and_evaluation_criteria": "yes", "theoretical_claims": "yes", "experimental_designs_or_analyses": "yes", "relation_to_broader_scientific_literature": "This paper connects to the subfield of variance-reduction and pre-conditioning. Notably, it proposes an empirically well-performing variance-reduction algorithm, while prior variance-reduction algorithms don't perform well in practice.", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "This paper has concrete results and shows significant novelty. Although variance-reduction algorithms have better theoretical convergence rate due to the slightly strong assumption (mean-square smoothness), they generally perform poorly on empirical tasks and lose to state-of-art optimizers such as Adam and its variants. Notably, MARS outperforms those algorithms on training language models, not only in sample complexity (number of tokens) but also in oracle complexity (as measured by wall-clock time). The novelty is thus showing that variance-reduction algorithms can be effective in practice. In addition, to my knowledge this is the first algorithm that incorporate preconditioning technique with variance-reduction algorithms. On the theory side, this paper also provides a theoretical justification of why properly choosing the control parameter $\\gamma$ for the noise correction can improve the performance, and it also provides convergence guarantee of MARS. For these reasons, I'd recommend acceptance.", "other_comments_or_suggestions": "- I think the name MARS-Shampoo is a little misleading. First, eq (3.24) is only true if we turn off the pre-conditioning aggregate of $L_t,R_t$ in Shampoo and use the instantaneous gradients. Also, the update rule of $m_t \\mapsto UV^\\top$ (either directly with SVD or approximately with Newton-Schulz) is more aligned with muon than shampoo.\n\n- I believe the theorems prove convergence of MARS without clipping $c_t$ (which is not exactly the same algorithm in the pseudo-code)?\n\nTypos:\n- line 98: to reaches -> to reach\n- 238 (right): line 2 -> line 5\n- above Sec. 3.2.3: \"we present a lemma...\" is the lemma missing?\n- 836: $\\nabla f$ -> $\\nabla F$\n- eq (C.2): missing a )", "questions_for_authors": "1. Experiment results in Figure 1 suggests that muon has worse performance compared to AdamW and Lion; while in muon's report [1], muon outperformance these two optimizers. Could the authors briefly discuss this experiment difference? Is it due to different tasks with different scales? Also, have the authors tried the performance of MARS-Shampoo with different $\\gamma_t$'s other than $1-\\mu$ (which approximately recovers muon)?\n\n2. I find the equivalence between nesterov update and momentum with noise-corrected gradients very interesting (Lemma D.1). Do authors know prior works stating the same result? Or is this a new finding?\n\n3. How do Theorem B.4 and B.5 imply a convergence rate of $\\tilde O(T^{-1/3})$ for $\\mathbb{E}\\\\|\\nabla F(x_t)\\\\|^2$? I was trying to directly apply the smooth descent lemma $F(x_{t+1})-F(x_t) \\le \\langle \\nabla F(x_t), x_{t-1}-x_t\\rangle + \\frac{L}{2}\\\\|x_{t+1}-x_t\\\\|^2$, but then $x_{t+1}-x_t = -\\eta_t H_t^{-1}m_t$ and $(F(x_{t+1})-F(x_t))/\\eta_t \\le - \\langle \\nabla F(x_t), \\nabla F(x_t) - (F(x_t)-H_t^{-1}m_t)\\rangle + \\frac{L}{2\\eta_t}\\\\|x_{t+1}-x_t\\\\|^2$, which requires a bound on $\\mathbb{E}\\\\|\\nabla F(x_t) - H_t^{-1}m_t\\\\|^2$ instead of $\\mathbb{E}\\\\|\\nabla F(x_t) - m_t\\\\|^2$.\n\n4. I find the empirical choices of $\\gamma_t$, shown in Figure 14, very interesting. (As a side note, I find it a bit hard to read the colors in the plot as many curves overlap with each other.) The optimal value of $\\gamma=0.025$ is very close to the value of $1-\\mu$, in which case the momentum with correct gradient is roughly momentum with nesterov update and MARS-Adamw should be close to AdamW with nesterov update. Have the authors tested the empirical performance of nesterov AdamW (or Adan) and compared it with MARS? Also, could the authors further elaborate on $\\gamma\\approx 1-\\mu$? It makes me wonder whether variance-reduction is the right way to interpret MARS, since the noise-correction only plays a small role. Could it be possible that nesterov update (or optimistic update) is a better interpretation?\n\n5. Have the authors tried MARS without preconditioning? I think the comparison of MARS-AdamW and MARS-SGDM (which is presumably significantly worse) would help motivating that both variance-reduction and preconditioning are necessary to achieve a good performance.\n\n[1] Keller Jordan, \"Muon: An optimizer for hidden layers in neural networks\", https://kellerjordan.github.io/posts/muon/, 2024.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "sLP2fiz5Ff", "reviewer": "Reviewer_XvSa", "summary": "The paper presents an algorithm template for variance reduction that includes common optimizers such as AdamW and Lion as a special case. A convergence analysis is provided the smooth case, with an assumption on the pre-conditioning matrix. Extensive experiments are conducted on language modeling and vision tasks.", "questions": "", "limitations": "", "rating": 5, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your careful review and constructive feedback. Your comments on both the theoretical and experimental aspects of our work have helped us further strengthen our contributions. We have carefully considered your suggestions and provide detailed responses below.\n\n---\n**Q1**: Parameters $\\beta_1$ and $\\beta_2$ may vary with time. I do not see this stated anywhere in the algorithm boxes and accompanying discussions\n\n**A1**: Thank you for your question. Our use of time-varying $\\beta$ follows the line of theoretical work on Adam, where $\\beta$ is dependent on $\\eta_t$. We will add a note that, for theoretical analysis, $\\beta_{1, t}$ and $\\beta_{2, t}$ are required. In practice, these parameters are typically chosen as constants.\n\n---\n**Q2**: (1) In Lemma C.4, the problem constant D is not formally introduced until Lemma D.2. (2) The learning rate sequence is set, but is also accompanied by an upper bound\n\n**A2**: Thank you for your observation. (1) We will revise the definition of constant D into Lemma C.4 and Theorem B.5. (2) We will revise the upper bounds on $\\eta_t$ in Theorems B.4 and B.5 by expressing them as conditions on the parameter $s$, avoiding any implicit assumptions on problem-specific constants. Specifically, we will require $s\\ge 8L^3/\\rho^3$ in Theorem B.4 and $s\\ge \\max(8L^3/\\rho^3, 64\\lambda^3)$ in Theorem B.5.\n\n---\n**Q3**: I believe there should be some discussion about Assumption B.3. Is this common in the literature? \n\n**A3**: Thank you for your question. Assumption B.3 is not a strong convexity assumption, but rather a condition for ensuring computational stability. This assumption is common in the literature--for example, Assumption 3 in Huang et al., (2021) and AS.3 in [1]. It is more often used implicitly, as in [2], where $H_t = \\eta / (\\sqrt{\\hat{v}_t} + \\lambda)$ in Algorithm 1, and Theorem 4.1 dependent on $\\lambda$. The same holds for Algorithm 2 in [3] and many others. We will include a more detailed discussion of this assumption and its role in related work in our revised manuscript.\n\n[1] Wang et al., Stochastic Quasi-Newton Methods for Nonconvex Stochastic Optimization  \n[2] Li et al., Convergence of Adam Under Relaxed Assumptions  \n[3] Zhuang et al., Adabelief Optimizer: Adapting Stepsizes by the Belief in Observed Gradients\n\n---\n**Q4**: While the proofs are well-structured and easy to follow logically, the computations are quite heavy. If the authors are able, a bit of readability (using colors, etc) would be helpful\n\n**A4**: Thank you for the suggestion. While the rebuttal format limits manuscript changes, we will revise the proofs in the final version to follow the style of the analysis in Section 5.4.4 of Bach (2024).\n\n---\n**Q5**: It would be helpful for the authors to provide FLOPs in the rebuttal\n\n**A5**: Thank you for raising this point. We now provide a detailed FLOPs analysis for our Transformer experiments (with GPT-2 small as an example).\n\nFor a Transformer with batch size $B$, sequence length $s$, $l$ layers, hidden size $h$, and vocabulary size $V$, the total number of trainable parameters is $N = h(12lh + 13l + V)$ (approximately $1.24 \\times 10^8 $). The forward pass requires $N_F=2Bsh(12lh + 2sl + V) $ FLOPs (around $1.40 \\times 10^{14}$), and the backward pass roughly doubles that.\n\nFor optimizer updates:\n- AdamW: 11N FLOPs (≈ $1.36 \\times 10^9$)\n- MARS-AdamW (approx): 17N FLOPs (≈ $2.10 \\times 10^9$), accounting for $c_t$ clipping\n- Muon: 11N + 12h^2l(30h + 20) FLOPs (≈ $1.96 \\times 10^{12}$)\n\nWe have added Figure 2 (https://anonymous.4open.science/r/M-E6D7/F.pdf) to illustrate loss vs. FLOPs across different optimizers.\n\n---\n**Q6**: Related literature on control variates\n\n**A6**: We appreciate the reviewer’s insight and agree with the relevance to stochastic simulation and control variates literature. We will discuss the connection with the suggested papers in the final version. We would like to point out that, in methods like STORM, which are not strictly unbiased, the update does not follow the classical control variate form $\\mathbb{E}||X-\\gamma(Y-E[Y])||^2$. Instead, they control estimation error recursively through differences across consecutive iterates. When casting the STORM-style estimator in the simplified form, we define $Y=\\nabla f(x_{t+1},\\xi_{t+1})-\\nabla f(x_t,\\xi_{t+1})$, and the estimation error we aim to reduce becomes $\\mathbb{E}||X+\\gamma Y-E[Y]||^2 $(as shown on Page 4 of our manuscript, full formulation at Line 1172). Unlike standard control variates, the optimal $\\gamma^*$ here cannot be directly represented in terms of $Y-E[Y]$, making the connection less explicit.\n\n---\n**Q7**: The reader should be exposed to the schedule for $γ_t$ earlier\n\n**A7**:  Thank you for your suggestion. In the final version, we will include a more detailed discussion of the schedule for $\\gamma_t$ in the methodology section.", "claims_and_evidence": "The theoretical claims have accompanying proofs and the experimental results are accompanied by both code and data.", "methods_and_evaluation_criteria": "Overall, they do, as the most important evaluations are the optimization performance on respectably-sized deep learning models.", "theoretical_claims": "I have read the claims and the proofs briefly. While theory is generally not valued as much in deep learning optimization, I do believe there is some sloppiness in the presentation of the results that could be improved.\n\n1. It appears from the theorem statements that they analyze an algorithm for which the parameters $\\beta_1$ and $\\beta_2$ may vary with time. I do not see this stated anywhere in the algorithm boxes and accompanying discussions, and the reader should not have to discover this. Similarly, in Lemma C.4, the problem constant $D$ is not formally introduced until Lemma D.2. Finally, the learning rate sequence is set, but is also accompanied by an upper bound. In general, this would be invalid as it placed an assumption on the problem constants, but because there is a free parameter $s$, the condition that $\\eta_t \\leq ...$ and $s \\geq 1$ in Theorem B.4 and B.5 should be converted into a condition of the form $s \\geq ...$. \n2. I believe there should be some discussion about Assumption B.3. Is this common in the literature? Considering the extreme case in which one could compute or even approximate the Hessian as the pre-conditioner, this assumption is essentially assuming that the objective is strongly convex, which disagrees completely with the deep learning objective.\n3. On a more superficial note, I believe the authors meant to write $F(\\mathbf{x}_{T+1})$ in line 916. While the proofs are well-structured and easy to follow logically, the computations are quite heavy. If the authors are able, a bit of readability (using colors, etc) would be helpful. This does not factor into my score, but I still request this revision. See the analysis of SAGA in 5.4.4 of [Bach (2024)](https://www.di.ens.fr/~fbach/ltfp_book.pdf) for a reader-friendly Lyapunov stability analysis.\n\nThis said, it is laudable that the authors provided this analysis, which is likely a further step than similar papers in the area.", "experimental_designs_or_analyses": "The designs are overall valid, and I appreciate that wall clock time was used in some of the experiments. What is slightly concerning is that there seems to be no experiments/discussions surrounding FLOPs. I believe it is very important in this field to be aware of the exact computational cost at this scale. The authors seem to have access to industrial-level computing resources, making wall time a sometimes unhelpful performance measure for researchers, as the relative impact of additional FLOPs per iteration might change between differing resources (older GPUs, etc). This is my main concern overall, so it would be helpful for the authors to provide these numbers in the rebuttal.", "relation_to_broader_scientific_literature": "Personally, I would be excited to finally see a variance reduction method actually become performant in deep learning, and much of the relevant literature in convex optimization is cited. That being said, I wish the authors were a bit more generous toward the stochastic simulation literature. The discussion on page 4 starting from line 206 seems to suggest that tunable variance reduction is a recent idea developed in deep learning optimization. Rather, nearly all of the (at least unbiased) variance-reduced gradient estimators are control variates, and setting $\\gamma = 1$ has been a purposeful choice since SVRG. See [Asmussen & Glynn (2007)](https://link.springer.com/book/10.1007/978-0-387-69033-9), Chapter V Section 2 on control variates or [Mehta et. al. (2024)](https://openreview.net/pdf?id=TTrzgEZt9s) Eq. (6) for a simpler derivation than the one on page 4 of the submitted paper.\n\nOn a similar note, I feel as though the authors do not seriously discuss the schedule for $\\gamma_t$ until Appx. F. Because this is one of the main components of the algorithm, the reader should be exposed to it earlier.", "essential_references_not_discussed": "None that I am aware of beyond the discussion above.", "other_strengths_and_weaknesses": "The paper is generally well-written, and Section 3.2 is especially nice. My main concerns are some of the theoretical statements, experiments with FLOP count, and some discussions that would benefit the main text. I would raise my score to a 5 if all were addressed.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "None other than those brought up earlier.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "IOmBY3xgfq", "reviewer": "Reviewer_ZHwA", "summary": "Training large deep neural networks is a challenging endeavor due to their massive scale and non-convexity. \nThis is true even more so for large language models where the current state-of-the-art method AdamW. \nThe paper propose incorporating a form of variance reduction, ala STORM, and combining it with a family of preconditioned adaptive gradient methods. \nNotably, much work over the past decade has tried to incorporate variance reduction (which has exhibited great success in training convex ML models) with Deep Learning.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your careful evaluation of our work. We appreciate your recognition of our contributions and your constructive questions. We have carefully considered your concerns and we provide detailed responses to your questions below.\n\n--- \n**Q1**: First sentence in the abstract: Training deep neural networks—and more recently, large models demands efficient and scalable optimizers.\" I'm assuming you meant to include language after large?\n\n**A1**: Thank you for pointing this out. We intentionally omitted the word “language” to highlight the broader applicability of our approach beyond language models. As demonstrated in Appendix E.4, we include experiments on computer vision tasks to support the general applicability of the MARS optimizer. Nevertheless, we appreciate your feedback and will revise the sentence to avoid confusion.\n\n---\n\n**Q2**: The essential related work is complete, but to give a more thorough picture of the optimization in ML literature, you should cite the following two recent papers.\n\n**A2**: We appreciate the suggestion and will incorporate these references into our related work section. [1] introduces SVRN, a method that combines variance reduction with second-order optimization. It achieves faster convergence rates for smooth, strongly convex problems while enabling large-batch parallelization, demonstrating advantages over first-order methods and subsampled Newton approaches. [2] proposes a suite of sketching-based preconditioned algorithms (including SVRG, SAGA, and Katyusha variants) for ill-conditioned convex problems. Their work introduces \"quadratic regularity\" to guarantee linear convergence even with infrequent preconditioner updates, outperforming tuned baselines on regression tasks.\n\n[1] Dereziński, M. (2022). Stochastic variance-reduced newton: Accelerating finite-sum minimization with large batches.\n[2] Frangella, Z., Rathore, P., Zhao, S., & Udell, M. (2024). Promise: Preconditioned stochastic optimization methods by incorporating scalable curvature estimates.\n\n---\n\n**Q3**: On novelty: The main optimizer proposed here MARS-AdamW is very close to VRAdam, Algorithm 2 in Li et al. 2024, in that the variance reduction mechanisms are almost the same, except that MARS-AdamW has some scaling factors not present in VRAdam. Nevertheless, I was wondering if the authors could comment on the relation between the two algorithms? Also, I'm very aware that although seemingly small, such things can have a big impact on practical performance. So I don't mind that much that it's somewhat similar to existing approaches.\n\n**A3**: Thank you for bringing our attention to VRAdam. While MARS-AdamW and VRAdam share the high-level goal of integrating variance reduction into adaptive methods, there are several key differences. \n\nFirst, unlike VRAdam (and SuperAdam by Huang et al., 2021), MARS defines the second-order momentum $v_t$ as the exponential moving average of the squared norm of the corrected gradient $c_t$, rather than that of the raw stochastic gradient. This modification is critical for maintaining the correct update scale in each coordinate direction. \n\nAdditionally, MARS introduces two further components: a tunable scaling factor $\\gamma_t$ and gradient clipping on $c_t$, both of which are essential for stability and performance, especially in deep learning scenarios. \n\nMoreover, we introduces weight decay in both our implementation, as well as the theoretical analysis, while VRAdam does not.\n\nLastly, MARS is designed as a flexible framework that extends beyond AdamW. Its design allows easy adaptation to other optimizers such as Lion and Shampoo, highlighting its broader applicability.\n\nWe hope this clarifies the distinctions and contributions of our approach. We will add the above discussion in our revised manuscript.\n\n---\n\n**Q4**: Open Source Implementation: The algorithms seem quite effective. Do you plan on providing an open-source implementation if the paper is accepted? I think it could be useful for researchers working in this area.\n\n**A4**: Thank you for your kind words and support for open research. We have included our implementation in the supplementary materials and plan to open-source the code upon acceptance. We also welcome feedback from the community to further improve it.", "claims_and_evidence": "**Experimental**\n\n- The experimental evaluation in the paper is very strong. The proposed MARS optimizer with AdamW or Lion both outperform tuned AdamW (the gold-standard for LLMs) on some standard large-language training tasks. The framework is also tested on some more classical tasks in the supplement on some benchmark computer visions tasks like CIFAR-10 and CIFAR-100, where the proposed optimizer also shows good performance.\n- The supplement also contains comparisons across various hyperparameter values, so the paper is quite thorough. \n\n**Theory**\n- In Appendix B.2, the authors show under standard assumptions, that MARS AdamW converges to an $\\epsilon$-approximate stationary point after $\\tilde O(\\epsilon^{-3})$ iterations, which matches the known lower bound (Arjevani et al. '23) in the literature up to log-factors. \nThus, in terms of theory, you can't really reasonably ask for a much better result without having to impose stricter assumptions. \n\n**Overall**\n\n- I think the results in the paper are compelling, and the proposed framework has promise to prove training of LLMs and large deep nets. \nSo I recommend acceptance. \n\n**References**\nArjevani, Yossi, Yair Carmon, John C. Duchi, Dylan J. Foster, Nathan Srebro, and Blake Woodworth. \"Lower bounds for non-convex stochastic optimization.\" Mathematical Programming 199, no. 1 (2023): 165-214.", "methods_and_evaluation_criteria": "The datasets, models, and optimizers used in the paper for evaluation are all logical choices for the setting of the paper. AdamW is the most widely used optimizer in LLM training so it provides a natural baseline. Muon is a relatively recent optimizer, but has been shown to outperform AdamW on certain LLM tasks. So the baselines employed in the paper are sound.", "theoretical_claims": "The proofs I read in detail seem correct to me. I read  the full proofs for the main theorems, Theorems B.4 and B.5, and I skimmed the proofs of the supporting lemmas.", "experimental_designs_or_analyses": "The experiment design and the conclusions drawn seem fine to me. The paper uses standard benchmark models and does a thorough hyperparameter evaluation for their method and the competing methods.", "relation_to_broader_scientific_literature": "The paper introduces an optimization framework based on variance-reduction that yields algorithms with the potential to make training of large models more efficient. This is of course a timely question given importance of LLMs and the fact they are so expensive to train. \nI also find it interesting that the paper gives real practical evidence that variance reduction can be helpful in Deep Learning. Since Defazio and Bottou (2019), the conventional wisdom is that variance reduction just isn't useful in Deep Learning. \nBased on the results of this paper, this no longer seems to be the case.\nSo, I think the paper's results are quite interesting relative to the existing literature.", "essential_references_not_discussed": "The paper has done a thorough job reviewing variance reduction and adaptive gradient methods in the context of deep learning. All the major works are either mentioned in the main paper or in Appendix A of the supplement.", "other_strengths_and_weaknesses": "**Additional Strengths**\n\n- The paper is well-written and very easy to follow.\n\n**Weaknesses**\n\n- Overall the paper, does not have any significant weaknesses in my view. It might be a little lacking in novelty in certain places (see my question in the box below), but in my view that takes a backseat to the fact that framework shows good practical performance and has good theoretical support.", "other_comments_or_suggestions": "**Typos**\n-  First sentence in the abstract:\n\n\"Training deep neural networks—and more recently, large models demands efficient and scalable optimizers.\" \n\nI'm assuming you meant to include language after large?\n\n**Suggestions**\n\nThe essential related work is complete, but to give a more thorough picture of the optimization in ML literature, you should cite the following two recent papers:\n\n- Derezinski, Michal. \"Stochastic Variance-Reduced Newton: Accelerating Finite-Sum Minimization with Large Batches.\" In OPT 2023: Optimization for Machine Learning.\n- Frangella, Zachary, Pratik Rathore, Shipu Zhao, and Madeleine Udell. \"Promise: Preconditioned stochastic optimization methods by incorporating scalable curvature estimates.\" Journal of Machine Learning Research 25, no. 346 (2024): 1-57.\n\nThese papers combine variance reduction with preconditioning in the convex setting, and show it makes for a very strong combination, both in theory and in practice. So, I think these works are very much in alignment with the papers message.", "questions_for_authors": "**On novelty**\n- The main optimizer proposed here MARS-AdamW is very close to VRAdam, Algorithm 2 in Li et al. 2024, in that the variance reduction mechanisms are almost the same, except that MARS-AdamW has some scaling factors not present in VRAdam. Nevertheless, I was wondering if the authors could comment on the relation between the two algorithms?  Also, I'm very aware that although seemingly small, such things can have a big impact on practical performance. So I don't mind that much that its somewhat similar to existing approaches. \n\n**Open source implementation**\n- The algorithms seem quite effective. Do you plan on providing an open-source implementation if the paper is accepted? I think it could be useful for researchers working in this area.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper proposes applying a scaled version of the STORM momentum that has an extra control parameter before preconditioning, and shows that this captures a wide range of optimizers with momentum in practice while achieving better results with the extra knob. The reviewers all appreciate how the paper ties together existing works on variance reduction methods, though noting a relatively small change and wondering how it changes previously negative results. Some concerns on the experimental setting remain, for instance Muon being worse than Adam can be specific to this setting given the emerging literature that contradicts this conclusion (e.g., Liu et al., 2025). It also omits certain baselines due to computational reasons (e.g., the Shampoo variant). While limitations on novelty and experiments remain, the authors find the authors' rebuttal sufficient. Following the reviewers' recommendation, I'd advise the authors to reduce the main content on repeating known optimizers and focus on the experiments to answer the more important questions raised above (e.g., what are the conditions for variance reduction to work).
Decision: Accept (poster)