Gradient Multi-Normalization for Efficient LLM Training

We thank the reviewer for the detailed feedback and for acknowledging the theoretical grounding, practical motivation, and empirical performance of SinkGD. We address the concerns below.

---

• On ablation analysis.

  (1) Hyperparameter sensitivity: SinkGD has very few hyperparameters. The number of Sinkhorn iterations is ablated in the paper. The learning rate is directly transferred from Apollo [2] and used across all model sizes without tuning. This demonstrates robustness to hyperparameter choices.

  (2) Norm choice: we emphasize that SinkGD is defined by row- and column-wise norm normalization. All other common matrix norms (e.g., spectral norm, $l_\infty$ norm) should correspond to other baselines (SWAN, AdamW) already included in our comparisons. Thus, norm choice is implicitly ablated through baseline comparisons.

  (3) Stability: SinkGD is stateless and does not accumulate second moments. This avoids instability issues observed in Adam due to large update norms from moment ratios [1]. SinkGD also avoids instability from low-rank approximations, which are observed in [4] and [6]. Moreover In Table 2, our 1 B model trained with SinkGD matches the performance of 7B models trained with Galore/Apollo/Adam across training stages, showing no signs of instability. In contrast, instability of Adam under the same setup is reported in [2].

---

• On baseline coverage.

  We acknowledge that including more baselines is beneficial. Due to time constraints during the rebuttal, we were unable to complete benchmarking of Lion/Sophia under our setup. However, our paper targets memory-efficient LLM pretraining, and we follow the baseline setup of Galore [3], Apollo [2], SWAN [5], Fira [4], and GWT [6], none of which compare with Lion/Sophia as they operate under different memory regimes.

  Additionally, our paper compares with 8 strong baselines, including SWAN, Apollo, and Fira, which are more recent than Lion and Sophia. We believe this provides a fair and comprehensive evaluation within the scope of memory-efficient stateless training.

---

• On generalization beyond C4 pretraining.

  Our focus is on scalable stateless training. Prior to our work, even demonstrating fast and stable stateless pretraining at the scale considered in this paper had not been established. We believe showing that a simple method (SinkGD) derived from MNGD can match or outperform Adam (with similar computational cost) in this setting is already a meaningful contribution.

  Moreover, our setup closely aligns with recent work on memory-efficient optimization such as Galore [3], Apollo [2], and SWAN [5], which also focus primarily on pretraining, where memory constraints and efficiency are critical.

---

• On novelty and contribution.

  We respectfully disagree with the assessment that our contribution is merely a synthesis of existing components. If we look at the literature

  • signSGD [7] (and arguably, reduced Adam) uses the sign operator, a known mathematical primitive.
  • SWAN [5] and Muon use whitening/orthogonalization operators from prior literature such as computer vision architectures.
  • Galore [3] uses SVD/subspace projections, well-known in optimization.

  However, these optimizers are all considered significant contributions. This tells us that we must carefully distinguish the projection function itself used in updates, vs the underlying optimization principle that derived and propose to use those projections, that leads to practical gains that are not predicted by prior works.

  In our case:

  1. To the best of our knowledge, no prior work has proposed using Sinkhorn iterations as a gradient projector; and no prior work has made sufficient hint that could trivially predict the practical gain of such operations.
  2. Our work is the first to propose performing steepest descent under an ensemble of multiple norms. This meta-viewpoint enables the synthesis of new optimizers beyond SinkGD and opens up a rich design space. While SWAN touches on this from the perspective of transformer dynamics, it does not explicitly uncover or formalize this principle.

---

References
[1] Molybog et al., "A theory on Adam instability in large-scale machine learning", arXiv:2304.09871 (2023)
[2] Zhu et al., "Apollo: SGD-like memory, AdamW-level performance", arXiv:2412.05270 (2024)
[3] Zhao et al., "Galore: Memory-efficient LLM training by gradient low-rank projection", arXiv:2403.03507 (2024)
[4] Chen et al., "Fira: Can We Achieve Full-rank Training of LLMs Under Low-rank Constraint?", arXiv:2410.01623 (2024)
[5] Ma et al., "SWAN: SGD with Normalization and Whitening Enables Stateless LLM Training", arXiv:2412.13148 (2024)
[6] Wen et al., "Breaking memory limits: Gradient wavelet transform enhances LLMs training", arXiv:2501.07237 (2025)
[7] Bernstein et al., "signSGD: Compressed optimization for non-convex problems", ICML (2018)

Dear Reviewer JbUr,

We appreciate the time and care you devoted to evaluating our manuscript.

We’ve submitted our formal rebuttal and would welcome your thoughts on whether it satisfactorily addresses the issues you highlighted. Any additional remarks you can share would be greatly valued.

Should you require further clarification before the discussion window closes on August 6, we remain at your disposal.

With sincere thanks, The Authors

We thank the reviewer for the positive assessment and for highlighting the clarity of our presentation, the theoretical rigor of our framework, and the empirical strength of SinkGD.

We address the concerns below.

---

• On generalization beyond C4 and pretraining.

  Our focus is indeed on scalable stateless for pretraining. Pretraining from scratch provides a clean and controlled setting to isolate the optimizer behavior. We also acknowledge the importance of fine-tuning, nevertheless, they introduces more confounding factors compared to pretraining (e.g., model quality, task-specific dynamics, optimizers used to pretrain the model) that masks optimizer effects. That said, we include preliminary fine-tuning results on GLUE datasets (comparing SinkGD to Galore, rank=8):

  Dataset	QNLI	MRPC	RTE	STS-B
  Galore (rank=8)	92.20	92.01	79.78	90.82
  SinkGD	92.05	92.38	79.78	91.19

  SinkGD performs comparably while maintaining lower memory usage. We will include these results in the revision and plan to explore larger-scale and cross-domain applications in future work.

---

• On hyperparameter tuning of baselines.

  For Galore and Apollo, we directly use the optimal configurations reported in their respective papers. This is rigorous for the following reasons:

  1. We use the exact same dataset, model architecture, and training setup.
  2. We use the same training and evaluation codebases (Apollo’s repo inherits from Galore).
  3. Both papers report that their methods are relatively insensitive to hyperparameter choices such as learning rate.
  4. For SinkGD, we reuse the global learning rate and local scaling from Apollo without additional tuning.
  5. Regarding the rank parameter: our method is stateless, which corresponds to rank zero in Galore/Apollo. Thus, the comparison is conservative and favors the baselines (as both Galore and Apollo uses $r>1$).

---

• On the use of Adam for embedding and output layers.

  We follow the Galore and Apollo setup, which also retain Adam for these layers. After submission, we found that SinkGD can be applied to all layers. However, for embedding and LM head layers, we observed that reintroducing momentum is important to match performance of SinkGD for Attention/MLP + Adam for embedding. We will include this observation in the revision.

---

• On the intuition behind multi-norm normalization.

  Two perspectives:

  1. From prior work: SWAN, the first stateless optimizer, uses multiple norms motivated by transformer dynamics. Our framework generalizes this into a principled multi-norm formulation.
  2. From optimization theory: the choice of norm affects convergence. For example, recent work (Xie et al., 2024) shows that Adam exploits $\ell_\infty$ geometry to reduce the empirical smoothness constant. This suggests that norm choice should reflect the loss landscape. Our approach can be viewed as an ensemble of norms, providing regularization and robustness across diverse geometries.

---

• On the computational cost of Algorithm 3.

  Both SinkGD and Adam have the same complexity ($\mathcal{O}(m^2)$ for square matrices). In practice, especially under distributed training, the cost of Sinkhorn iterations is negligible compared to forward/backward passes. For example, 5 iterations of Sinkhorn involve roughly $5d^2$ flops, while training cost scales with batch size $\times d^2$. Since batch sizes are large in practice, the relative overhead of SinkGD diminishes at scale.

---

We also appreciate the reviewer’s suggestions and will incorporate the clarifications and additional results in the revised version.

---

Reference

1, Xie, Shuo, Mohamad Amin Mohamadi, and Zhiyuan Li. "Adam Exploits $\ell_\infty$-geometry of Loss Landscape via Coordinate-wise Adaptivity." arXiv preprint arXiv:2410.08198 (2024).)

We thank the reviewer for the thoughtful review and for highlighting the theoretical foundation of our framework, the memory and efficiency gains of SinkGD, and its strong empirical performance in LLM pretraining.

We address the concerns below.

---

• On the theoretical limitations of convergence guarantees.

  We acknowledge that our current analysis only guarantees convergence to a fixed point and does not establish global optimality of the multi-norm problem (Eq.4 of the paper). While we do not yet have a complete theoretical explanation, our empirical results suggest that a fixed point solution provide sufficiently adaptive regularization to enable effective training. We believe this opens a promising direction for future theoretical work.

---

• On broader applicability and scaling beyond 1.3B.

  Our goal is to enable scalable stateless training. Prior to our work, even demonstrating fast and stable stateless pretraining at the 1.3B scale had not been established. We believe showing that a simple method like SinkGD can match or outperform Adam (with similar throughput) in this setting is already a meaningful contribution.

  Furthermore, in Table 2, we show that a 1.3B model trained with SinkGD matches the performance of a 7B model trained with memory-efficient or quantized Adam variants from the literature. This suggests that SinkGD is competitive even when compared to significantly larger models trained with more complex optimizers.

  Regarding other domains: we include preliminary fine-tuning results on GLUE datasets (comparing SinkGD to Galore, rank=8):

  Dataset	QNLI	MRPC	RTE	STS-B
  Galore (rank=8)	92.20	92.01	79.78	90.82
  SinkGD	92.05	92.38	79.78	91.19

  SinkGD performs comparably while maintaining lower memory usage. We will include these results in the revision and plan to explore larger-scale and cross-domain applications in future work.

---

• On the novelty of SinkGD beyond SWAN.

  While SinkGD builds on the idea of stateless gradient normalization, our contribution is not merely an incremental improvement over SWAN. The SWAN paper motivates its design from a specific observation of transformer dynamics in a toy setting. Their work mainly says: the first order dynamics of transformer implies you should use row-normalization; and its second order dynamics implies you should use orthogonalization (ie normed GD under spectral norm). However, SWAN does not discuss why these two operations should be used jointly. In contrast, we address this issue and generalize the core principles behind SWAN into a unified framework of multi-normalization, which recovers SWAN as a special case and opens up a broader design space for stateless optimizers. This new development is not covered/implied in the SWAN paper, we believe this generalization is conceptually significant and will motivate further development of efficient optimizers for large-scale models.

---

We appreciate the reviewer’s suggestions and will incorporate the clarifications and additional results in the revised version.

We thank the reviewer for the constructive feedback and for recognising the novelty of our multi-norm perspective, the memory efficiency of SinkGD, and its strong empirical performance in LLM pretraining.

The concerns raised are mostly due to scope clarification and implementation details. We address them below.

---

• SinkGD results beyond pretraining.

  Our focus is on scalable stateless training for LLMs, where pretraining from scratch provides a clean and controlled setting to isolate the optimizer behaviour. We also acknowledge the importance of fine-tuning, nevertheless, they introduces more confounding factors compared to pretraining (e.g., model quality, task-specific dynamics, optimizers used to pretrain the model) that obscure optimizer effects. That said, we include preliminary fine-tuning results on GLUE datasets (comparing SinkGD to Galore, rank=8):

  Dataset	QNLI	MRPC	RTE	STS-B
  Galore (rank=8)	92.20	92.01	79.78	90.82
  SinkGD	92.05	92.38	79.78	91.19

  SinkGD performs comparably while maintaining lower memory usage.

---

• BF16 was chosen for consistency with prior work.

  BF16 is the standard precision used in Galore, Apollo, and SWAN. It reflects realistic memory constraints in large-scale training. To address the reviewer’s concern, we ran additional experiments in mixed-precision (FP32) on a 350M model (batch size: 250K tokens, 30K steps):

  Steps	BF16 SinkGD Val PPL	Mixed-Precision SinkGD Val PPL
  10K	20.34	20.64
  20K	17.10	17.38
  30K	16.40	15.66

We observed that under mixed precision, the convergence is indeed stronger; however, BF16 SinkGD offers stronger early-mid stage convergence. (A side note: this observation might imply a hybrid approach for efficient training.)

---

• Assumption 3.3 beyond Euclidean norm (l2 norm)

Can authors provide other norm examples except for the Euclidean norm that satisfies Assumption 3.3?

Sure. For example, the common choices for optimization, such as $l_\infty$ norm (used in Adam), and spectral norm (matrix case, used in SWAN). We will clarify this in revision.

---

• SinkGD achieves higher throughput due to statelessness and parallelism.

  While SinkGD performs iterative normalization, it avoids maintaining and updating first/second moment states (as in AdamW), as well as subsquent element-wise operations. Additionally, our implementation shards the Sinkhorn computation across devices. Each local rank processes a subset of rows/columns, enabling efficient parallelism. This explains the higher raw throughput observed in Table 3.

---

We appreciate the reviewer’s suggestions and will incorporate the clarifications and additional results in the revised version.