MARS: Unleashing the Power of Variance Reduction for Training Large Models

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.

---

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?

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.

---

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.

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.

[1] Dereziński, M. (2022). Stochastic variance-reduced newton: Accelerating finite-sum minimization with large batches. [2] Frangella, Z., Rathore, P., Zhao, S., & Udell, M. (2024). Promise: Preconditioned stochastic optimization methods by incorporating scalable curvature estimates.

---

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.

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.

First, 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.

Additionally, 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.

Moreover, we introduces weight decay in both our implementation, as well as the theoretical analysis, while VRAdam does not.

Lastly, 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.

We hope this clarifies the distinctions and contributions of our approach. We will add the above discussion in our revised manuscript.

---

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.

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.

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.

---

Q1: I'm not sure if 'unified' is the right word here, because variance reduction is orthogonal to preconditioning

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.

To 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.

---

Q2: Are there characteristics of settings in which variance reduction should help

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.

We 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).

---

Q3: The paper would benefit by bringing some of the results to the main paper and deferring details about the instantiations

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.

---

Q4: Typos

A4: Thank you for spotting these typos. We have corrected them and will ensure they are reflected in the final version.

---

Q5: What was the reasoning behind the choice in not including MARS-Shampoo for the language model pretraining experiments?

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.

---

Q6: Learning rate stability of MARS-AdamW

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.

---

Q7: Sweep across batch size

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.

---

Q8: Compare to methods like AdEMAMix

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.

Thank you for your detailed and thoughtful review. Below, we provide detailed responses to each of your comments.

---

Q1:The name MARS-Shampoo

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.

[1] Bernstein et al., Old optimizer, new norm: An anthology

---

Q2:The theorems prove convergence of MARS without clipping

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.

In 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.

[1] Zhang et al., Why gradient clipping accelerates training: A theoretical justification for adaptivity.

---

Q3:Typos

A3: Thank you, we have corrected them.

---

Q4:Experimental difference between Muon's report and Figure 1 of our paper

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.

---

Q5:Is Lemma D.1 a new finding.

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.

---

Q6: How do Theorem B.4 and B.5 imply a convergence rate

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.

---

Q7: Empirical performance of Nesterov AdamW (or Adan). Also, could the authors further elaborate on γ≈1-μ?

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).

For 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:

(a) Comparing exact vs. approximate variance reduction (Figure 5). and

(b) We explore γ values not strictly tied to 1-μ (e.g., γ=0.025 vs. μ=0.95).

We 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.

---

Q8: MARS-SGDM

A8: We have now included a comparison between MARS-AdamW and MARS-SGDM in Figure 4 in the link above.

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.

---

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

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.

---

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

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.

---

Q3: I believe there should be some discussion about Assumption B.3. Is this common in the literature?

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.

[1] Wang et al., Stochastic Quasi-Newton Methods for Nonconvex Stochastic Optimization
[2] Li et al., Convergence of Adam Under Relaxed Assumptions
[3] Zhuang et al., Adabelief Optimizer: Adapting Stepsizes by the Belief in Observed Gradients

---

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

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).

---

Q5: It would be helpful for the authors to provide FLOPs in the rebuttal

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).

For 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.

For optimizer updates:

• AdamW: 11N FLOPs (≈ $1.36 \times 10^9$)
• MARS-AdamW (approx): 17N FLOPs (≈ $2.10 \times 10^9$), accounting for $c_t$ clipping
• Muon: 11N + 12h^2l(30h + 20) FLOPs (≈ $1.96 \times 10^{12}$)

We have added Figure 2 (https://anonymous.4open.science/r/M-E6D7/F.pdf) to illustrate loss vs. FLOPs across different optimizers.

---

Q6: Related literature on control variates

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.

---

Q7: The reader should be exposed to the schedule for $γ_t$ earlier

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.