Title: Gradient Multi-Normalization for Efficient LLM Training
Venue: NeurIPS 2025
Accepted: yes
Keywords: LLM

Optimizer

Gradient

Normalization
Abstract: Training large language models (LLMs) commonly relies on adaptive optimizers such as Adam (Kingma & Ba 2015), which accelerate convergence through moment estimates but incur substantial memory overhead. Recent stateless approaches such as SWAN (Ma et al., 2024) have shown that appropriate preprocessing of instantaneous gradient matrices can match the performance of adaptive methods without storing optimizer states. Building on this insight, we introduce \emph{gradient multi-normalization}, a principled framework for designing stateless optimizers that normalize gradients with respect to multiple norms simultaneously. Whereas standard first-order methods can be viewed as gradient normalization under a single norm (Bernstein & Newhouse, 2024), our formulation generalizes this perspective to a multi-norm setting. We derive an efficient alternating scheme that enforces these normalization constraints and show that our procedure can produce, up to an arbitrary precision, a fixed-point of the problem. This unifies and extends prior stateless optimizers, showing that SWAN arises as a specific instance with particular norm choices. Leveraging this principle, we develop SinkGD, a lightweight matrix optimizer that retains the memory footprint of SGD while substantially reducing computation relative to whitening-based methods. On the memory-efficient LLaMA training benchmark (Zhao et al., 2024), SinkGD achieves state-of-the-art performance, reaching the same evaluation perplexity as Adam using only 40\% of the training tokens.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a new family of stateless optimizers, termed Multi-Normalized Gradient Descent (MNGD), to improve the memory efficiency of large language model (LLM) training. The method generalizes prior work such as SWAN by normalizing stochastic gradients with respect to multiple matrix norms through a proposed multi-step alternating projection algorithm. The authors further present a specific instance, Sinkhorn GD (SinkGD), which uses the Sinkhorn algorithm to implement the multi-normalization efficiently. Experimental results on LLaMA pretraining tasks (up to 1B parameters) show that SinkGD achieves comparable or better perplexity than Adam while significantly reducing memory consumption and training time.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths：\nThe paper presents a novel family of stateless optimizers, Multi-Normalized Gradient Descents (MNGDs), which generalizes prior work like SWAN by normalizing gradients with respect to multiple norms. The proposed SinkGD variant leverages the Sinkhorn algorithm for efficient and scalable normalization, and is theoretically justified through convergence analysis. Empirical results demonstrate improvements in memory efficiency and throughput in large language model (LLM) pretraining tasks, suggesting practical utility.\n\nWeaknesses:\n1.Limited experimental scope: All experiments are restricted to LLaMA pretraining on C4. There is no downstream task or broader benchmark evaluation to verify robustness or generality.\n2.Weak ablation analysis: The only ablation varies the number of Sinkhorn iterations. Other key factors, such as norm selection, convergence stability, or sensitivity to optimizer hyperparameters, are not analyzed.\n3.Insufficient baselines: Several strong recent optimizers (e.g., Sophia, Lion variants, or other stateless optimizers) are missing from the comparison, which limits the strength of the empirical claims.\n4.Incremental novelty: While the proposed formulation is cohesive and practically motivated, it mainly combines existing components (e.g., Sinkhorn projection). The contribution may be viewed more as a practical synthesis than a fundamentally new optimization paradigm.\n\nQUESTIONS:\n1.Broader Benchmark Evaluation: The paper only provides experiments on LLaMA pretraining with the C4 dataset. Have the authors considered evaluating SinkGD on downstream tasks such as language modeling benchmarks (e.g., WikiText, OpenBookQA, etc.) or real-world applications to demonstrate generalization?\n2.Baseline Coverage: While comparisons to Adam, SWAN, and a few others are presented, important recent optimizers (e.g., Sophia, Lion variants, AdaFactor) are missing. Could the authors justify these omissions or provide further experiments?\n3.Ablation Scope: The only ablation study concerns the number of Sinkhorn iterations. Could the authors provide more analysis on the role of norm choices, projection quality, or convergence stability across different architectures and model sizes?\n4.Scaling and Practicality: While the paper claims that SinkGD is more memory efficient and computationally attractive, the actual implementation details (e.g., compatibility with distributed training frameworks or wall-clock speed) are limited. Can the authors clarify whether SinkGD integrates smoothly into existing LLM training pipelines at scale?\n\nLIMITATIONS:\nWhile the paper proposes a promising stateless optimizer with theoretical grounding, several limitations remain. \n1. The current experiments are confined to LLaMA pretraining on the C4 dataset. Without evaluation on downstream tasks, transfer benchmarks, or domain-shifted settings, it remains unclear how well the method generalizes beyond its pretraining context.\n2. The ablation analysis is limited to varying the number of Sinkhorn iterations. Other key aspects—such as the choice of norm, stability across training regimes, or sensitivity to optimizer hyperparameters—are not investigated, leaving gaps in understanding the method’s behavior.\n3.  The empirical section omits comparisons with several competitive recent optimizers, including Sophia, Lion variants, and other stateless or memory-efficient methods. This weakens the strength of the performance claims made in the paper.\n4. While the formulation is well-integrated and practically motivated, the core components—such as multi-norm projections and the use of Sinkhorn iterations—are largely drawn from existing techniques. As such, the contribution may be interpreted more as an effective synthesis rather than a fundamentally novel optimization approach.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-16"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper propose a new framework for stateless optimizers named as Gradient Multi-Normalization. Previous approaches such as SWAN could be interpreted as specializations of Gradient Multi-Normalization. The authors also provide theoretical convergence guarantee for iterative multi-normalization. Based on their framework, the authors propose SinkGD, a new memory-efficient optimizer. SinkGD achieves on-par or better performance compared with previous methods on LLM pretraining tasks, with improved memory costs and convergence speed.\n\nSTRENGTHS AND WEAKNESSES:\n### Clarity:\nThe paper is well-written, with clear presentation of theoretical derivations and detailed experimental evaluations.\n\n### Quality:\nThe theoretical insights behind the gradient multi-normalization framework are well-explained and supported with rigorous proofs.  The experiments are comprehensive and demonstrates the superiority of proposed method.\n\n### Significance\n\n#### Strengths\nIn general, I think this paper makes a solid improvement in stateless optimizer design. The MNGD framework provides a new, insightful perspective on stateless optimizers, unifying and generalizing previous work like SWAN. The author demonstrates the strong empirical performance of their optimizer in memory efficiency and training speed-up.\n\n#### Weaknesses\n- The experiments are only conducted on the C4 dataset. Generalization to other tasks such as (parameter-efficient) finetuning is underexplored.\n- The authors did not mention how hyper-parameters in baseline methods (such as rank $r$ in Galore/Fira) are tuned. This could potentially lead to unfair comparisons in Table 1.\n\n### Originality\n\nResults obtained in this paper are new and original.\n\nOther comments:\n- Line 208: typo \"SR-Sinhkorn\" \n- The supplementary material appears to be wrongly uploaded appendix of another paper\n\nQUESTIONS:\n- What is the rationale for retaining Adam for the embedding and output layers? How much additional overhead will it bring? What would happen if SinkGD is used for all layers?\n\n- What is the intuition behind normalizing over multiple norms in MNGD?\n\n- While stateless optimization reduces memory costs compared with Adam-based methods, how much additional computational cost does the iterative sub-process in Algorithm 3 incur?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper tackles the high-memory cost of adaptive optimizers such as Adam in LLM training and proposes a stateless alternative that remains competitive in speed and final accuracy. It does so by framing gradient preprocessing as a multi-normalization problem, proving that simple alternating projections can enforce several norm constraints simultaneously, and instantiating this idea in a lightweight optimizer dubbed SinkGD. Experiments on 350 M- and 1.3 B-parameter LLaMA variants show that SinkGD matches or surpasses Adam’s perplexity while cutting optimizer-state memory to (roughly) SGD levels and delivering 2-3× token-efficiency gains.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths:**\n\n1.Theoretical Foundation: The paper provides a principled framework (MNGD) and proves key properties of the proposed optimizer. In particular, it shows that the alternating multi-normalization procedure can converge to a fixed-point solution with arbitrary precision. This theoretical result, along with connections to known algorithms, lends credibility to the method’s soundness.\n\n2.Memory and Efficiency Gains: By eliminating optimizer state, SinkGD matches the memory footprint of plain SGD. The method significantly reduces computation compared to SWAN by removing the expensive whitening step. This yields Adam-level step time complexity instead of the higher complexity introduced by SWAN’s whitening, making the approach more scalable without sacrificing accuracy.\n\n**Weakness:**\n\n1.Limitations in Theorem: While the authors prove convergence to a fixed point, they acknowledge that the algorithm is not guaranteed to *optimally* solve the original optimization problem in a theoretical sense. The convergence analysis assumes certain conditions and does not provide guarantees for non-convex objectives beyond reaching a fixed point. This leaves a gap in understanding whether the procedure always leads to the best solution or just a stable one. \n\n2.Scope of Empirical Validation: The experiments, though thorough on LLaMA pre-training, are limited to language model pretraining tasks up to 1.3B parameters. The method’s efficacy on *larger* models or other domains (e.g. fine-tuning, different architectures) remains unproven. \n\n3.Not Enough Contribution: The main algorithm SinkGD is essentially an “improved SWAN” that relaxes SWAN’s constraints for efficiency. Thus,  the work is building on an existing concept (stateless gradient normalization) more than creating a wholly new paradigm.\n\nQUESTIONS:\nThe experiments focus on LLaMA language model pre-training (with model sizes from 60M up to 1.3B parameters). It would help if the authors discuss the broader applicability of their optimizer. For example, can the proposed stateless optimizer scale to even larger models (tens or hundreds of billions of parameters) without issues? Would it be effective for other training scenarios like fine-tuning, different architectures, or large-batch training beyond the C4 dataset?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes an optimization framework which processes the gradient/momentum by applying multiple normalizations. Under this framework, they develop a practical algorithm named SinkGD, which iteratively performs row normalization and column normalization on the momentum without introducing heavy computation overhead. The algorithm is shown to be convergent based on the theory of Sinkhorn algorithm. Numerical experiments show that SinkGD achieves significantly faster convergence than AdamW (2.4-2.8 $\\times$) on the task of pretraining LLMs in C4 dataset.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n\n* The perspective of steepest descent under multiple norm measurements is novel and unifies existing algorithms such as Muon and SWAN.\n* SinkGD avoids storing second moment term as in AdamW, which is memory-efficient. The iterative normalization operation seems cheap based on the throughput in Table 3.\n* SinkGD attains superior convergence under the LLM pretraining tasks and outperforms AdamW by a large margin, which partially justifies the effectiveness of the multi-normalization step. The algorithm can be of general interest given its memory efficiency and fast convergence.\n\n**Weaknesses**\n* The experiments focus on pretraining stage. To broaden its applicability, I suggest to conduct additional experiments of SinkGD under finetuning stage, where the memory budget is usually more limited.\n* The experiment results are based on pure BF16 precision. It is suggested to provide the numerical results under mixed-precision training scheme, which is often used in LLM's pretraining.\n\nQUESTIONS:\n* Can authors provide other norm examples except for the Euclidean norm that satisfies Assumption 3.3?\n* In Table 3, the raw throughput of SinkGD is higher than AdamW, which is confusing to me. AdamW only needs to perform element-wise operations twice, while SinkGD requires iteratively normalizing the gradient for multiple steps.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "CLARITY": 3, "IMPACT": 4, "ORIGINALITY": 3, "DATE": "2025-07-05"}