COSMOS: A Hybrid Adaptive Optimizer for Efficient Training of Large Language Models
Authors: Liming Liu
Zhenghao Xu
Zixuan Zhang
Hao Kang
Zichong Li
Chen Liang
Weizhu Chen
Tuo Zhao
Venue: ICLR 2026 Poster (2026)
Keywords: Efficient Optimizer for LLMs
Preconditioning
Muon
SOAP
Adam
Primary area: optimization
TL;DR: Not All Eigensubspaces of the Gradient are Created Equal

Abstract: Large Language Models (LLMs) have demonstrated remarkable success across various domains, yet their optimization remains a significant challenge due to the complex and high-dimensional loss landscapes they inhabit. While adaptive optimizers such as AdamW are widely used, they suffer from critical limitations, including an inability to capture interdependencies between coordinates and high memory consumption. Subsequent research, exemplified by SOAP, attempts to better capture coordinate interdependence but incurs greater memory overhead, limiting scalability for massive LLMs. An alternative approach aims to reduce memory consumption through low-dimensional projection, but these methods lose the gradient information in the residual space, resulting in less effective optimization. In this paper, we propose COSMOS, a novel hybrid optimizer that leverages the varying importance of eigensubspaces in the gradient matrix to achieve memory efficiency without compromising optimization performance. The design of COSMOS is motivated by our empirical insights and practical considerations. Specifically, COSMOS applies SOAP to the leading eigensubspace, which captures the primary optimization dynamics, and MUON to the remaining eigensubspace, which is less critical but computationally expensive to handle with SOAP. This hybrid strategy significantly reduces memory consumption while maintaining robust optimization performance, making it particularly suitable for massive LLMs. Numerical experiments on various datasets and transformer architectures are provided to demonstrate the effectiveness of COSMOS.

Published as a conference paper at ICLR 2026
COSMOS: A HYBRID ADAPTIVE OPTIMIZER FOR EFFICIENT TRAINING OF LARGE LANGUAGE MODELS
Liming Liu Georgia Tech
Zhenghao Xu Georgia Tech
Zixuan Zhang Georgia Tech
Hao Kang Georgia Tech
Zichong Li Georgia Tech
Chen Liang Microsoft
Weizhu Chen Microsoft
Tuo Zhao
Georgia Tech
ABSTRACT
Large Language Models (LLMs) have demonstrated remarkable success across various domains, yet their optimization remains a significant challenge due to the complex and high-dimensional loss landscapes they inhabit. While adaptive op- timizers such as A are widely used, they suffer from critical limitations, including inability to capture interdependencies between coordinates and high memory consumption. Subsequent research, exemplified by SOAP, attempts to better capture coordinate interdependence but incurs greater memory overhead, limiting scalability for massive L.LMs. An alternative approach aims to reduce memory consumption through oow-dimensional projection, but these methods lose the gradient information in the residual space, resulting in less effective optimiza- tion. In this paper, propose COSMOS, novel hybrid optimizer that lever- ages the varying importance of eigensubspaces in the gradient matrix to achieve memory efficiency without compromising optimization performance. The design of COSMOS is motivated by our empirical insights and practical considerations. Specifically, COSMOS applies SOAP to the leading eigensubspace, which cap- tures the primary optimization dynamics, and MUON to the remaining eigensub- space, which is less critical but computationally expensive to handle with SOAP. This hybrid strategy significantly reduces memory consumption while maintain- ing robust optimization performance, making it particularly suitable for massive LLMs. Numerical experiments on various datasets and transformer architectures are provided to demonstrate the effectiveness of COSMOS.
INTRODUCTION
The optimization of Large Language Models (LLMs) is fundamental to their success, enabling these models to achieve ofate-of-the-art performance across diverse tasks. However, the non-convex loss landscapes inherent to LLMMs, which can contain hundreds of billions or even trillions of parame- ters (Achiam et al. 2023), present significant optimization challenges. Adaptive optimizers, such as Adam (Kingmal 2014) and its variants AdamW (Loshchilov) 2017), have emerged as de facto stan- dards due to their ability to dynamically adjust learning rates based on the second moment of the gradient. Despite their widespread adoption, these methods suffer from two critical limitations that impede their effectiveness and scalability the context of increasingly large and complex LLMs:
(1) Adam and its variants have limitations in adaptive learning rates. By adjusting learning rates independently for each parameter, the method reduces computational complexity but may not fully capture parameter interdependencies. In complex architectures of LLMs, this independent approach can lead to suboptimal parameter updates (Zhang et al. 2024a).
(II) Another limitation of Adam and its variants lies in the substantial memory requirement for storing per-parameter adaptive learning rates and gradient statistics. As LLM sizes increase, memory consumption becomes prohibitively large, impeding scalability.
Correspondence 11iu606, zhenghaoxu, zzhang31 05, tourshaolgatech.ee))

---

Published as a conference paper at ICLR 2026
To address the limitations of Adam and its variants, researchers have pursued two main approaches. The first approach, exemplified by algorithms such as Shampoo (Gupta et al. 2018) and the more recent SOAP (Vyas al. 2024), employs sophisticated techniques to capture curvature information and parameter interdependencies. These methods utilize rotational matrices, derived through (ap- proximate) singular value decomposition (SVD) of the gradient matrix, to provide more compre- hensive representation of the loss landscape's geometry. This approach allows for better approx- imation of the full preconditioning matrix, enabling the capture of inter-coordinate dependencies. However, the improved capability of representing parameter interactions comes at the cost of sub- stantial computational and memory overhead, rendering these algorithms challenging to implement for large-scale LLMs, where memory efficiency is crucial.
The second approach focuses on reducing memory consumption through various approximation techniques. Algorithms such AdaFactor (Shazeer and Stern 2018) and Adam-mini (Zhang et al.) 2024b) aim to decrease memory usage by approximating the second moment of the gradient ma- trix. Adam-mini employs component-specific approach, averaging second moments neuron wise for certain layers. Meanwhile, AdaFactor utilizes a rank-l approximation of the second moments. While these methods reduce memory cost, their approximations oversimplify the structure of the gradient matrix's 's second order moments, compromising optimization performance. The trade-off between memory efficiency and the preservation of gradient statisties remains a crucial challenge.
More recent approaches, such as Gal.ore (Zhao et al. 2024a) and MUON (Jordan et al. 2024 have attempted strike balance between computational complexity, memory consumption, and optimization performance in LLM training. GaLore, which can be viewed as a memory-efficient variant of SOAP, approximates the first and second moments of the gradient matrix in the leading eigensubspace. While it effectively reduces memory consumption, Liang et al. (2024) find that its effectiveness diminishes for sequence lengths exceeding 256. MUON, essentially an approximation of Shampoo based on Newton-Schulz transformation proposed in Bernstein and Newhouse (2024), aims to decrease computational complexity. However, MUON only estimates the eigensubspaces based on the gradient on one batch, rather than capturing the comprehensive distribution of gradients across the entire optimization process.
In this paper, we propose COSMOS, a novel hybrid optimizer that addresses the limitations of ex- isting methods by exploiting the varying importance eigensubspaces the gradient matrix. Our approach decomposes the gradient into two parts: projection onto the leading eigensubspace and projection onto the remaining eigensubspace. The leading eigensubspace captures the most signifi- cant directions of change in the gradient, typically corresponding the most important optimization dynamics. For this part, we apply a S-AP-like optimization strategy. However, by crucially restrict- ing SOAP the leading eigensubspace, COSMOS only needs to maintain the projection matrix and the second-order moment within this small subspace, thereby retaining SOAP's ability to capture parameter interdependencies while substantially lowering its memory cost. The remaining eigen- subspace, while less critical, still significantly influences optimization performance. To address this, we employ MUON as more efficient alternative to SOAP for this digh-dimensional space. Such a hybrid approach allows COSMOS to maintain optimization effectiveness while significantly reduc- ing memory requirements compared to SOAP, potentially enabling the training of larger LLMs or the use of increased batch sizes.
We highlight the key contributions of this paper as follows: (1) We propose a novel hybrid opti- mization strategy. This leads us to develop the COSMOS algorithm, which synergizes the strengths of SOAP and MUON by decomposing the gradient matrix into eigensubspaces of varying impor- tance. (2) COSMOS achieves significant memory consumption reduction compared to the SOAP algorithm, while achieving equally or better optimization performance.
RELATED WORK
The optimization of LLMs has seen significant advancements in recent years, with various ap- proaches aimed at improving efficiency and performance. This section discusses key related works in adaptive optimization, memory-efficin techniques, and specialized algorithms for LLMs.
Cuordinate-wise adaptive optimizers: Adam (Kingma| 2014) and AdamW (Loshchilov) 2017) have become standards in deep learning optimization due to their ability to dynamically adjust
2

---

Published as a conference paper at ICLR 2026
learning rates based the first and second moments of the gradients. However, these methods treat parameters independently, failing to capture interdependencies between coordinates. This limitation can lead to suboptimal updates, especially in the complex architectures of LLMs. Other adaptive optimizers such as Lion (Chen et al., 2023), Sophia (Liu et al. 2023), and Adafactor (Shazeer and Stern 2018 Zhai et al.; (022)) have shown comparable performance AdamW in L.M pretraining but have not significantly surpassed it, suggesting the need for non-diagonal preconditioners.
Second-Order Optimizers: Researchers have explored second-order optimization techniques for training large models. These methods can be broadly categorized into Hessian-free approaches and Hessian estimation methods. Hessian-free methods, such as those proposed by MMatens (2010) and Martens and Grosse (2015), optimize without explicitly computing the Hessian matrix. On the other hand, Hessian estimation methods maintain an efficient approximation of the Hessian for neural networks. Notable examples include KFAC (Martens and Grosse 2015). Shampoo (Gupta et al., 2018) and SOAP (Vyas et al. 2024).
Shampoo and Its Variants: gorithm, is motivated by the online learning algorithm Adagrad (Duchi et al. 2011). Shampoo also employs layer-wise Kronecker-factoree
Shampoo (Gupta et al.) 2018), another second-order optimization al-
preconditioner. A recent distributed implementation of Shampoo (Shi et al. 2023) won an optimization efficiency benchmark (Dahl et al. (2023), highlight-
ing the practical utility of second-order methods in deep learning. Other works (Anil et al.) 2020
Peirson et al.
2022: Lin et al. 2024 Wang et al. 2024 Zhao et al.
2024b have proposed various
strategies improve Shampoo's scalability.
SOOAP: SOAP algorithm (Vyas et al. 2024) establishes a formal connection between Shampoo and Adafactor. SOAP is equivalent to running Adafactor in the eigenbasis of Shampoo's preconditioner, leading to a simpler and computationally efficient algorithm. By continually updating the running average of the second moment in the current (slowly changing) coordinate basis, SOAP mitigates the performance degradation associated with less frequent cigendecomposition computations. SOAP has shown significant improvements over AdamW in per-token efficiency.
Memory-efficient optimizers: As LLM sizes increase, memory efficiency becomes crucial. Several approaches have been proposed to reduce the memory footprint optimizers:
Adafactor and Adam-mini: Shazeer and Stern (2018) use low-rank approximation of the sec- ond moments reduce memory consumption. has been widely used in LLMs due to memory efficiency. ZZhan et al. (2024b) achieve comparable performance than Adam with 50% smaller memory footprint. It reduces memory by carefully partitioning parameters into blocks and assigning single learning rate to each block based on the Hessian structure of neural networks.
GaLore: Zhhao et al. (2024a) reduce Adam's 's memory footprint by maintaining momentum a low-rank subspace derived from the singular value decomposition (SVD) of the gradients. However, its effectiveness diminishes for sequence lengths exceeding 256, as shown in Liang et al. (2024).
MUON: The MUON optimizer (Jordan et al. 2024) can be viewed an efficient approximation of Shampoo. employs a Newton-Schulz transformation to approximately implement the Kronccker- factored preconditioner. While computationally more complex than Adam, MUON only adds minor overhead to the overall training time due to efficient parallelization of matrix operations.
These advancements highlight the efforts to improve the training efficiency and performance of LLMs. However, each approach comes with its own trade-offs in terms of computational complexity, memory requirements, and performance. Our work builds upon these insights develop hybrid approach that aims to balance these factors effectively, combining the strengths of different methods to achieve both memory efficiency and robust optimization performance for massive LLMs.
3 COSMOS: A HYBRID ADAPTIVE OPTIMIZER
We present a novel hybrid optimizer COSMOS in Algorithm T which can achieve memory effi- ciency without compromising performance for training LLMs. Without loss of generality, we use m and n to denote the numbers of rows and columns in a m by n matrix, and we assume m >> n. For simplicity, the following notations:
3

---

Published as a conference paper at ICLR 2026
Matrix Sign Operator: Given matrix X E R=xx and its reduced-SVD X UDVT, where D E R"** is a diagonal matrix containing all singular values of X, and E Rm*x and E Raxx are left and right singular vector matrices, respectively. We define
MatSgn(X)) UV
Newton Schulz (NS) transformation: Given matrix Xo E IRmx", where ||ollle S 1, we define
NS5 X5.
where Xs is obtained by X = aX, + b + X X2 for A: = 0, 1.... 4 with 3.4445, -.77750 and c 2.0315. Bernstein and Newhouse (2024) first mentioned this
transformation to approximate the matrix sign operator without specifying the coefficient. Pordan al. (022)) later used an ad-hoc gradient based approach to find the set of coefficients here.
Normalization operator: NORM(X) = / l|., where || || denotes the Frobenius norm. The normalization operator used to normalize the output of the NS transformation.
Gram-Sc SchnmSchidt procedure: (R(X).
Algorithm COSMOS for an m x layer W. Per layer, we maintain four matrices: U E Rnxr S E grxr, V R and M _
input Learning rate 7). combination weight % projection rank r < n, momentum parameters (31, perturbation parameter c. For simplicity, we omit the initialization.
1: for 1= O,...
2: Sample batch Mr
3:
G, < VwdM, (W,)
4:
Mc + -1MM11 + (1 B1
5:
U. < QR(BgUi-11S-11+((
8.)6 G,U,-))
6:
S. 18 (1- B, G.)U,
7:
V + B2Ve-1 +
@ (GU)
M,li((1 81)
8:
A, ((Vi +c) (5)
U.
Mr- M,U.U
9:
B, < NORM NS5
| M.U.U, lIF
10:
G.+ A 7 m
11: W141 + W. NORM(G.) Vm
12: end for
Design principle The design of COSMOS is guided by simple principle: instead of maintain- ing SOAP's full second-moment matrix- which is memory-rrohibitive- track its dominant eigenspace and operate in a low-dimensional subspace.
In the SOAP algorithm, the exponential moving average (EMA) of the second moment is
H. = B-H1-1 (1 G2) Gt,
where G, is the stochastic gradient. Because H, E RRxx is dense, storing is infeasible for large n. COSMOS avoids this by maintaining (i) an orthonormal basis U, E RRar for the leading eigenspace of H, and (ii) projected second-moment matrix S, grxr with
S, U. H,U,
Assume at step that U,-1 spans the dominant eigenspace and S -i Ap- proximating H.-1 by rank-r surrogate, UU- and substituting into Equation 1pp yields
H, B2 U1-1S--UU;
82) G. Gr.
We then update the basis via one-step power iteration:
U, QR (RH U..
4

---

Published as a conference paper at ICLR 2026
and refresh the projected second moment by
S, U, flUUG.
These two steps track the dominant eigenspace and its curvature information with (wr) memory.
Given Us Line 7 of Algorithm I maintains the EMA of the projected gradients V TR and Line 8 performs SOAP-like adaptive update within the subspace spanned by Us producing A, after projecting back to the full parameter space. This is the SOAP component COSMOS.
To complement the low-rank update, COSMOS applies MUON-inspinn preconditioner the orthogonal complement of U.. Writing the orthogonal projector as P.- U,U, Line forms
|MPP. ))
M,Pt
M MUUU
B, = NORM NS5
-NOM NORM NS5
(2)
| MUU,U
where NS5 applies directly to the residual momentum; no additional matrices are stored. Finally, Lines 10-11 combine the two components:
G A BBym,
W4+1 W. - NORM (G:) vm.
The normalization ensures the update has Frobenius norm O ((vmn) matching MUON's scaling. In sum, COSMOS adaptively preconditions the leading eigenspace as in SOAP while using MUON on the residual, achieving robust optimization with substantially reduced memory.
Algorithm 2 (One-side) SOAP
Algorithm MUON
input Learning rate 7/, momentum parameters (B1, B2). input Learning rate 7), momentum perturbation parameter e.
parameters μ..
1: for t = 0, ... do
1: for f = O,... do
2: Sample batch M,
2: Sample batch Mr
3: G, + VWŌM, (W))
3:
G, + Vw M. (W;)
4:
Me + 3 + - 31)B
4:
Mr + J&M + Gt
5:
Lt <- B2G, G,U,-1 + G3 G.U--1
5:
N₂ p.M + G2
6:
U. < (R(L,U,-1)
6:
B₁ t NS5((N/ll
7:
G + GiU,,
7:
W4+1 + We nBe vm
8:
VE 32Ve-1 + (1 (1) Gt)
8: end for
MUU/(1 - B4)
9: A. =
U,
(i, +,)(11 35)
W411 W. - -nAt
11: end for
Remark As can be seen, COSMOS only needs to maintain four matrices in the memory: M,
Rmxn, U. E Raxr, S. E Rr and V, E Rmxx. In sharp contrast, even one-sided SOAP
needs to maintain Mr E Rmxn L, E Raxn U, E guxa and V₁ E pm The resulting memory overhead is which is significantly larger than that COSMOS.
Remark Recall that the the computation complexity of the QR decomposition on a matrix of the shape n. is O(nr²) when rn, the low rank QR decomposition of B2U + (1 GT G,U,-. in COSMOS is actually very quick since r < n (and much quicker than that in SOAP, which is Onn²)). Therefore, unlike SOAP which needs consider the preconditioning frequency for performing QR decomposition, we can carry out QR decomposition at every step with virtually no overhead. In addition. PyTorch provides an efficient implementation of QR method, which is also used by SOAP. In Table we provide the comparison of wall-clock time per iteration to show that compared to MUON, COSMOS only incurs a very slight increase wall-clock time.
3.1 MEMORY USAGE COMPARISON
For comparison, list the memory usage of the optimization states in Adam, Adam-mini, SOAP, MUON and COSMOS for training transformer models in Table For simplicity, we assume that the attention weight matrices Wo, WK, Wy, Wo dxd and the ML.P weight matrices W, E Rdxid
5

---

Published as a conference paper at ICLR 2026
Table I: Memory usage of the optimization states in different algorithms for training transformers.
Adam
Adam-mini
SOAP
MUON
COSMOS
24d
12d
66d
12d
13d
and W2 pddddd Note that in practical LLMs, the dimensionalities of Wi and W2 might slightly vary. Moreover, we assume that the rank of the projection r = 0.05d for COSMOS.
We remark that Table T only compares the optimization states. In practice, however, besides the optimization states, the overall memory usage also includes the model weights and intermediate variables used in the forward and backward passes well additional memory overhead of I/O and computation. Therefore, we present more detailed and practical memory profiling for training LLaMA-IB IB model in Section
EXPERIMENTS
We evaluate the performance of COSMOS pre-training various sizes of LLMs, in comparison with baseline algorithms including Adam (Kingma 2014). Adam-mini (Zhang et al. 2024b), Ga- Lore (Zhao et al. 2024a), SOAP (Yass et 20024) and MUON (Jordan et al. 2024). Note that for SOAP, MUON and COSMOS, the embedding and output weights are trained by Adam.
Models and datasets. We train LLaMA-type models (Touvron al. 2023 the C4 dataset (Raffel et al. 2020), which is a colossal, cleaned version of Common Crawl's web crawl corpus for pre-taining. We conduct comprehensive experiments and ablation studies 130M models and demonstrate the token efficiency of COSMOS. We then scale up to 350M and IB models to show- case the memory efficiency and small computational overhead of COSMOS. Due limited compu- tational resources, experiments on these larger models are less comprehensive, while still capable of illustrating the efficacy of our method. We train for one epoch on portion of the C4 dataset, ranging from 5B to 26B tokens, and scaling with the model size according to the scaling law (Kaplan et al. 2020 We set the maximum sequence length as 1024 by default.
Besides LLaMA models and C4 dataset, we also conducted experiments with modded-NanoGPP (Jordan et al. 2224) on WineWeb (Penedo et al. 2024) and GPT-2 (Radford et al. (2099) on WikiTeett- (Merity et al. 2016) to evaluate the effectiveness of COSMOS across different settings.
103
4.1 COMPARISON ON LLAMA-130M
For LLLMA-130M, we train on a 5B subset of the C4 dataset. We compare COSMOS's validation loss with that of Adam, Adam-mini, Galore, MUON, and SOAP.
Hyperparameters. For each method, we tune the corresponding learning rate to obtain optimal performance. We select the rank r 64 for COSMOS and 256 for GaLore. To avoid multiple hyperparameter tuning, we set the discount factor in COSMOS as 7 17/1//. where 7oo is the learning rate of Adam for training the embedding and output weights in the implementation of COSMOS. Other hyperparameter choices follow Zhao et al. (022a)) and are provided in detail in Section A.I.I
Main results. We plot the validation loss curves in Figure As illustrated, COSMOS consistently outperforms MUON with better stability and is comparable to SOAP. This showcases that our hybrid approach leveraging the leading eigensubspace captures the most important information for efficient update, allowing COSMOS to achieve similar per-token efficiency as SOAP. In addition, all three methods outperform the vanilla Adam and are much better than Adam-mini and GaLore, validating that inter-coordinate dependence is crucial for efficient optimization. We report the final validation perplexity in Table 2
Ablation on learning rates. We experiment with different learning rates while keeping the rank r 64 and discount factor y = n/ no for COSMOS. As shown in the Table 3 COSMOS is not very sensitive the learning rate, and it achieves the best performance at 5e-4. As comparison, MUON is more sensitive to the learning rate, and it underperforms COSMOS across all learning rates.
Ablation on rank and discount factor. We also experiment with different ranks and discount factors y while keeping the learning rate as 5e-4 for COSMOS, and the results are summarized in Table As illustrated, COSMOS is not very sensitive and 7, and the best discount factor is

---

64 for Published as a conference paper at ICLR 2026
3.5
3.225
Adam
Adam Adam mini
mini
3.200
3.4
Muon
Galore
3.175
Soep
Muon
Soap
Cosinos
Voliiation Loos
3.150
3.3
3.2
3,075
3.1
3.050
3.025
3.0
1000
2000
3000
4000
5000
2000
3000
4000
5000
Update Steps
Update Steps
Figure 1:
Performance LLaMA-130ME COSMOS consistently outperforms baseline methods. In the right plot, we hide GaLore to better compare the performance of COSMOS with SOAP and MUON, as the curves are close in the left plot.
around 0.25 to 0.5 across different ranks. In practice, our choice Y = n/no falls in this range, so it serves as valid heuristic that prevents extra tuning of y. We also conduct experiments with 0.25 fixed while varying the learning rate. The results show that, when 7 = 0.25, COSMOS consistently outperforms Muon across different learning rates, as shown in Table We provide details in Section A Moreover, we observe that rank increases, COSMOS performs slightly worse and is more sensitive to the choice of 7. One possible explanation is that when the rank large, the top-r eigenvalues of M, contain some smaller values that are close to the remaining cigenvalues. For these cigenvalues, the one-step power iteration (Line in Algorithm [ cannot accurately approximate their corresponding cigensubspaces, leading to larger approximation errors and worse performance.
Table 2: Validation perplexity after training on C4 dataset. We train for 5000 steps on 130M and 350M models and 13000 steps on IB model. COSMOS achieves the best validation perplexity.
Size(Tokens)
130M(5B)
350M(10B)
1B(26B)
Adam
21.28
17.28
12.97
Adam-mini
21.78
18.03
GaLore
24.07
19.03
SOAP
20.59
16.32
MUON 20.69
16.49
12.57
COSMOS
20.54
16.21
12.46
Table 3: Validation perplexity under different learning rates. We set n/no and COSMOS. Our method outperforms MUON across all learning rates.
Ir
2e-4
5e-4
le-3
2e-3
MUON
21.72
20.75
20.69
26.74
COSMOS
21.17
20.54
20.62
21.00
Table 4: Validation perplexity under different learning rates. We set 7 0.25 and r 64 for COSMOS. Our method outperforms MUON across all learning rates.
2e-4
5e-4
le-3
2e-3
MUON
21.72
20.75
20.69
26.74
COSMOS
21.09
20.54
20.57
20.87
Effeet of normalization. COSMOS applies normalization step (Line in Algorithm D after the NS transformation compared to MUON (Algorithm 3) Empirically, we find that COSMOS also outperforms the normalized version of MUON (see Figure 7 in Section D.Ip. This implies that normalization is not the only driving foree behind COSMOS'S efficiency.
GaLore degradation on long sequences. our experiment, we observe that GaLore performs much worse than COSMOS and other baselines, including Adam. Such degradation is less sig-

---

Published as a conference paper at ICLR 2026
Table 5: Valid perplexity of COSMOS under different r and y. COSMOS is not very sensitive and 7 and consistently outperforms MUON (20.69) except for only one config (r = 128, y 1).
r
0.25
0.5
32
20.58
20.55
20.54
20.54
64
20.62
20.54
20.57
20.61
128
20.65
20.58
20.63
20.72
nificant the shorter sequences with length 256 (see Figure 10 in Section D.5| the setting adopted in the original GaLore paper (Zhao et al. 2024a). This observation GaLore degradation on long while similar phenomenon appearing in fine-tuning is
sequences aligns with Liang al. (2024), reported by [ann et al. (024)) In contrast, COSMOS consistently outperforms Adam from short to long sequences without suffering from degradation.
4.2 SCALING UP TO LLAMA-350M AND LLAMA-1B
To further illustrate the effiency of COSMOS, we scale up to larger models and more tokens. For LLaMA-350M, we train on 10B subset of the C4 dataset for 5000 steps. We compare COSMOS with Adam, Adam-mini, GaLore, MUON, and SOAP. For LLaMA-IB, we train on a 26B subset for 13000 steps. Given limited GPU resources, we compare COSMOS with Adam and MUON. Adam serves the standard baseline, while MUON achieves better performance than Adam while requir- ing less memory and little computational overhead. Although SOAP show superior performance among baselines, exceeds our available resources. More experiment details are provided in Sections A.1.2 and A.1.3
exclude it from our comparison as its complete training for the IB model
Main results of token efficiency. Figures 2a and 2b display the validation loss curves for the 350M and 1B models, and Table presents their final validation perplexities. COSMOS demonstrates superior performance compared all baselines across both model sizes, matching the results on the 130M model and showcasing consistent token efficiency across different model sizes.
Adam
Adam
2,80
Muon
3.05
Muon
Soap
Validitioo Loss
Cosmos
3.00
Cosmos
2.70
2.95
2.65
2.90
2.60
2.85
2.55
2.80
2000
3000
4000
5000
3000
5000
7000
9000
11000
13000
Update Stegs
Update Steps
(a) LLaMA-350M trained on C4 dataset.
(b) LLaMA-IB trained on C4 dataset.
Figure 2: Comparison of performance on LLaMA-350M and LLaMA- 1B trained on the C4 dataset.
Memory and computation time profiling. To illustrate the memory efficiency and small compu- tation overhead of COSMOS, we conduct profiling experiment on the IB model. We fix the batch size as 10 and the gradient accumulation steps as 25 for all methods and record the maximum GPU memory usage and time spent during the entire forward- backward propagation and optimizer up- date process for one iteration. As shown in Table COSMOS achieves much lower maximum GPU memory usage than Adam (6.8%) and SOAP (19.4%) ) with slightly more overhead compared to MUON. In terms of wall-clock time per iteration, COSMOS is comparable to MUON and is much better than SOAP. The fastest Adam method cannot achieve the same level of token efficiency as COSMOS. Therefore, COSMOS strikes a good balance between token and memorylcomputation overheads, achieving the best final perplexity much lower cost.
borward-backward
8

---

Published as a conference paper at ICLR 2026
To better compare the methods in practical setting, evaluate the maximal batch size and throughput of COSMOS and baselines on a single NVIDIA A100 GPU with 80GB memory. The input sequence length is set to 1024. As shown in Table 7 COSMOS is 10.8% faster than SOAP and comparable to MUON.
Table GPU memory usage and wall-clock time per iteration on model. We the batch size to be 10 for all methods. COSMOS has significantly less memory usage than SOAP and is comparable to memoor-efficient methods like MUON, without introducing much computation overhead.
Method
Memory
Wall-clock time
Adam
62.75 G
34.73 s
SOAP
72.58 G
39.51 8
MUON
58.25 G
35.56
COSMOS
58.47
35.75 S
Table 7: System performance single NVIDIA A100-80G GPU and corresponding throughput (number of samples processed per second on C4 dataset) of IB model. Max batch size is defined as the maximum number of samples that fit within the GPU's memory capacity. Throughput is reported as the number of samples the GPU processes per second (samples/s).
Method
Max batch size
Throughpu((sumpees)
Adam
13
7.24
SOAP
10
6.33
MUON
14
7.23
COSMOS
14
7.07
Wall-Clock time plot for LLaMA-1B. Based on the throughput we calculate in Table 7 we rescale the X-axis of Figure 2b to be wall-clock time and present the result in Figure 9 in Section D.4 D.4 Our
results indicate that, in terms of actual training time, COSMOS still outperforms MUON and the Adam baseline, demonstrating COSMOS's potential for efficient pretraining.
4.3 ADDITIONAL EXPERIMENTS ON OTHER SETTINGS
Most current works pretraining optimizers, due limitations in resources and time, focus single model architecture and a single dataset - for example, GaLore (Zhao et al. 2024a)LLaMM on C4), SOAP (Vyas al. (2024), OLMo (Groeneveld et al. 2024 C4) and Muon (Jordan al. (2024), NanoeIT on FineWeb), Our experiments on LLaMA with C4 are already aligned with these prior works, validating the performance of COSMOS. To demonstrate that COSMOS retains its advantages under other settings as well, however, we also conduct experiments in the following additional settings:
LLaMA 130M on FineWeb: To further assess the advantages of COSMOS over Muon across di- verse settings, we conduct experiments with LLaMA under substantially different configuration. Specifically, we train 130M-parameter model on FineWeb-10B, corresponding to approximately 4x the data size recommended by Chinchilla scaling. We apply a weight decay 0.1 all optimiz- ers. With the sequence length fixed at 1024, we increase the batch size to 4096. This setup features large batch size, large-scale data, and strong weight decay, making it closer industrial-scall pre- training conditions. The results are shown in Figure 3
on FineWeh: GPT Codified-NNnoGGP on FineWeb. Since Muon has already performed extensive hyperparameter tun- ing in this setting, we do not tune Muon again but use the provided reproducible log. For COSMOS, simply align the learning rate with that of Muon and follow other settings. We present the results in Figure See Section for an additional wall-clock time plot and the detailed configuration.
We also conduct experiments Muon's original setting, namely
GPT2 on Wikitext-103s We also trained GPT2- small and GPT2-mediug pare COSMOS with Adam and Muon. In this setting, COSMOS still outperforms Muon and Adam. We present the result in Figure See Section D.2 for the detailed configuration.
WikiText- 103 to com-

---

Published as a conference paper at ICLR 2026
3.35
3.30
Validation Loss
Soap
Cosmos os
3.05
1.00
Update Steps
Figure 3: Comparison of optimization performance on 130M LLaMA trained on FineWeb dataset, with large batch size and data size. Weight decay fixed at for all optimizers.
3.45
Muon
Muon
3.50
Cosmos
3.40
Cosmos
3.45
Validation
Validation adaaoo
3.40
3.30
3.25
3.35
3.20
3.30
Steps
Steps
(a) 124M Web.
GanoPP trained on
(b) 350M Modded-NanoGPP trained on Web.
Figure 4: Comparison of optimization performance on 124M and 350M Modded-NanoGPT trained on the Fine Web dataset. dataset.COSMOS consistently outperforms MUON.
3.2
3.25
Adam
Adam
Muon
3.00
Muon
3.0
Cosmos
2.75
Cosmos
2aiinng 1oss
Traninn Less
2.8
2.50
2.25
2.00
1.75
1.50
Update Steps
Update Steps
(a) GPT2-small trained on WikiText- 103.
(b) GPT2-medium trained WikiText- 103.
Figure 5: Comparison of COSMOS. MUON and Adam on WikiText- 103 using GPT2-small and GPT2-medium models. COSMOS consistently outperforms MUON and Adam.
5 CONCLUSION
We develop a hybrid adaptive optimizer, COSMOS, which leverages the varying importance of cigensubspaces in the gradient matrix to achieve token efficiency, memory efficiency, and high computation throughput simultaneously. By decomposing the gradient matrix into leading and re- maining eigensubspaces and applying SOAP-like and MUON-like updates to them correspondingly. COSMOS uses significantly less memory than SOAP while achieving equal better optimization performance, Comprehensive experiments show that COSMOS performs consistently well across different settings.
10

---

Published as a conference paper at ICLR 2026
REFERENCES
ACHIAM, J., ADLER, S., AGARWAT, S., AHMAD, L., AKKAYA, I., ALEMAN, F. L., ALMEIDA,
D., ALTENSCHMIDT, J., ALTMAN, S., ANADKAT, S. AL. (2023). Gpt-4 technical report. arXiv preprint arXiv:2303.0877.
ANIL, R., GUPTA, V., KOREN, T., REGAN, K. and SINGER, Y. (2020). Scalable second order optimization for deep learning. arXiv preprint arXiv:2002.0911..
BA, J., GROSSE, and MARTENS, J. (2017). Distributed second-order optimization using kronecker-factored approximations. In International Conference on Learning Representations.
BERNSTEIN, J. and NEWHOUSE, L. (2024). Modular duality in deep learning. arXiv preprint arXiv:2410.21265.
CHEN, X., LIANG, C., HUANG, D., REAL, E., WANG, K., LIU, Y., PHAM, H., DONG, X., LUONG, T.. HSIEH, C.-J. ET AL. (2023), Symbolic discovery of optimization algorithms. arXiv e-prints arXiv-2302.
DAILL, G. E., SCHNEIDER, F., NADO, Z., AGARWAL, N., SASTRY, C. S., HENNIG, P., MEDA- PATI, S., ESCHENHAGEN, R., KASIMBEG, P., SUO, D. ET AL. (2023). Benchmarking neural network training algorithms. arXiv preprint 2rXiv:2306.07179
DUCHI, J., HAZAN, E. and SINGER, Y. (2011). Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12.
ELFWING, S., UCHIBE, E. and DOYA, K. (2018), Sigmoid-weighted linear units for neural network function approximation in reinforcement learning. Neural networks, 107 3-11.
ESCHENHAGEN, R., IMMER, A.. TURNER, E., SCHNEIDER, F. and HENNIG, P. (2023). Kronecker- factored approximate curvature for modern neural network architectures. arXiv preprint arXiv:231 1.00636.
GAO, K., LIU, X., HUANG, Z., WANG, M., WANG, Z., XU, D. and YU, (2021). A trace- restricted frontceer-factored approximation to natural gradient. In Proceedings of the AAAI Con- ference on Artificial Intelligence, vol. 35.
GEORGE, T., LAURENT, C., BOUTHILLIER, X., BALLAS, N. and VINCENT, (2018). Fast approximate natural gradient descent in kronecker-factored eigenbasis. arXiv preprint 1rXiv:1800.00388..
GROENEVELD, D., BELTAGY, I., WALSH, P., BHAGIA, A., KINNEY, R., TAFJORD, O., JHA, A. H., IVISON, H., MAGNUSSON, I., WANG, Y. ET AL. (2024). Olmo: Accelerating the science of language models. arXiv preprint arXiv:2402.00838
GUPTA, V., KOREN, T. and SINGER, Y. (2018). Shampoo: Preconditioned stochastic tensor opti- mization. In International Conference Machine Learning. PMLR.
JORDAN, K., JIN, Y., BOZA, V., JIACHENG, Y., CECSISTA, F., NEWHOUSE,, L, and BERNSTEIN, J. (2024). Muon: An optimizer for hidden layers in neural networks.
https: //keleelrrrrr jordan github. 1 http//pssss.oooo posts/mann /mmon//
KAPLAN, J., MCCANDLISII S., HENIGHAAN, T.. BROWN, T. B., CHESS, B., CIILD,, R., GRAY, S., RADFORD, A., Wu, J. and AMODEI, D. (2020). Scaling laws for neural language models. arXiv preprint arXiv:2001.08311.
KINGMA, D. P. (2014). Adam: A method for stochastic optimization. arXiv preprint 142.v.14...6980.
LIANG, K.. LIU, B., CHEN, L. and LIU, Q. (2024). Memory-cfficient Ilm training with online subspace descent. arXiv preprint arXiv:2408. /2857.
LIN, W., DANGEL, F., ESCHENHAGEN, R., BAE, J., TURNER, R. E. and MAKHZANI, A. (2024). Can we remove the square-root in adaptive gradient methods? second-order perspective. arXiv preprint ar arXiv:2402.033. 03496.

---

Published as a conference paper at ICLR 2026
LIU, H., LI, Z., HALL, D., LIANG, P. and MA, T. (2023). Sophia: A scalable stochastic second- order optimizer for language model pre- traiiiinnng. arXiv preprint arXiv:2305.1 14342.
LIU, J. Su, J., YAO, X., JIANG, Z., LAI, G., DU, Y., QIN, Y., Xu, W. LU, E., YAN, J. ET AL. (2025). Muon is scalable for IIm training, arXiv preprini 2rXiv:2502.16982.
LIU, Y., ОТT, M., GOYAL, N., DU, J., JOSHI, M., CHEN, D., LEVY, O., LEWIS, M., ZETTLE- MOYER, L. and STOYANOV, V. (2019). Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv:1907.11666
LOSHCHILOV, I. (2017). Decoupled weight decay regularization. arXiv preprint arXiv1171.00101.
MARTENS, (2010). Deep learning via hessian-frre optimization. In Proceedings of the 27th International Conference International Conference on Machine Learning.
MARTENS, J., BA, J. and JOHNSON, M. (2018), Kronecker-factored curvature approximations for recurrent neural networks. In International Conference on Learning Representations.
MARTENS, and GROSSE, R. (2015). Optimizing neural networks with kronecker-factored ap- proximate curvature. arXiv preprint arXiv:1503.056%
MERITY, S., XIONG, C.,, BRADBURY, J. and SOCHER, R. (2016). Pointer sentinel mixture models. arXiv preprint arXhv:1609.08843
OSAWA, K., TSUJ, Y., UENO, Y., NARUSE, A., YOKOTA, R. and MATSUOKA, S. (2018). Large- scale distributed second-order optimization using froteceer-factoreed approximate curvature for deep convolutional neural networks. arXiv preprint arXiv: 120019.
PAN, R., LIU, X., DIAO, S.. Pl, R., ZHANG, J., HAN, and ZHANG, (2024). Lisa: Layer- wise importance sampling for memory-efficient large language model fine-tuning, arXiv preprint arXiv:2403.17919
PRIRSON, A., AMID, E., CHEN, Y., FEINBERG, V., WARMUTH, M. K. and ANIL, R. (2022). Fishy: Layerwise fisher approximation for higher-order neural network optimization. In Has Trained Yet? NeurlPS 2022 Workshop.
PENEDO, G., KYDLÍČEK, H., LOZHKOV, A., MITCHELL, M., RAFFEL, C. A., VON WERRA, L., WOLF, T. ET AL. (2024). The fineweb datasets: Decanting the web for the finest text data at scale. Advances Neural Information Processing Systems, 37 30810-0133000 1-8849..
PUIU, C. O. (2022a). Brand new k-facs: Speeding up k-fac with online decomposition updates. arXiv preprint arXiv:2210.0849.
PUIU, C. O. (2022b). Randomized k-faes: Speeding up k-fac with randomized numerical linear
algebra. In International Conference on Intelligent Data Engineering and Automated Learning.
RADFORD, A., WU, J., CHILD, R., LUAN, D., AMODEL, D., SUTSKEVER, I. (2019).
Language models are unsupervised multitask learners. OpenAl blog, 19.
RAFFEL, C., SHAZEER, N., ROBERTS, A., LEE, K., NARANG, S., MATENA, M., ZHOU, Y., LI, W. and LIU, P. J. (2020). Exploring the limits of transfer learning with unified text-to-text transformer. Journal of machine learning research, 21 1-67.
ROBERTS, J. D. (1980). Linear model reduction and solution of the algebraic riccati equation by use of the sign function. International Journal of Control, 32 677-687.
SHAH, I., POLLORENO, A. M., STRATOS, K., MONK, P., CHALUVARAIU, A., HOJEL, A., MA, A., THOMAS, A., TANWER, A., SHAH, D. J. ET AL. (2025). Practical efficiency of muon for pretraining. arXiv preprint arXiv:2505.02222.
SHAZEER, N. and STERN, M. (2018). Adafactor: Adaptive learning rates with sublinear memory
cost. arXiv preprint arXiv:
12

---

Published as a conference paper at ICLR 2026
SHI, H.-J. M., LEE, T.-H., IWASAKI, S.. GALLEGO-POSADA, J., LI, Z.. RANGADURAL, K., MUDIGERE, D. and RABBAT, M. (2023). A distributed data-parallel pytorch implementa- tion of the distributed shampoo optimizer for training neural networks at-scale. arXiv preprint arXiv:2309.06997.
SU, J. AUMEE,, M., LU, Y., PAN, S., Bo, W. and LIU, Y. (2024). Roformer: Enhanced trans- former with rotary position embedding. Neurocomputting. 568 127063.
TOUVRON, H., MARTIN, L., STONE, K., ALBERT, P., ALMAHAIRI, A., BABAEI, Y., BASH- LYKOV, N., BATRA, S., BIHARGAAA, P., BIIOSALE, S. ET AL. (2023). Llama 2: Open foundation and fine-tuned chat models. arXiv preprini urXiv:2307.09288
VYAS, N., MORWANI, D., ZHAO, R., SHAPIRA, I., BRANDFONBRENER, D., JANSON, L. and KAKADE, S. (2024). Soap: Improving and stabilizing shampoo using adam. arXiv preprint arXiv:2409.1332.
WANG, S., ZHOU, P., LI, J. and HUANG, H. (2024), 4-bit shampoo for memory-efficient network training. arXiv preprint arXiv:2405. 18/44.
ZHAL, X., KOLESNIKOV, A., HOULSBY, N. and BEYER, (2022). Scaling vision transformers. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition.
ZHANG, Y., CHEN, C., DING. T., LI, Z., SUN, R. and Luo, Z.-Q. (2024a). Why transformers need adam: hessian perspective. arXiv preprint arXiv 16788.
ZIIANG, Y., CHEN,, C., LI, Z., DING, T., WU, C., YE, Y., LUO, Z.-Q. and SUN, R. (2024b). Adam-mini: Use fewer learning rates to gain more. arXiv preprint arXiv:2406.1679 16793.
ZIao,, J., ZHANG, Z., CHEN, B., WANG, Z., ANANDKUMAR, A. and TIAN, Y. (2024a). Galore: Memory-efficient IIm training by gradient low-rank projection. arXiv preprint arXiv:2403.03507
ZHAO, R., MORWANI, D.,, BRANDFONRRENER, D., VYAS, N. and KAKADE, S. (2024b). Decon- structing what makes good optimizer for language models, arXiv preprint arXiv:440..0797.
13

---

Published as a conference paper at ICLR 2026
A EXPERIMENT DETAILS ON LLAMA MODELS
Many aspects of our setup such as models are the same in Zhao al. (2024a). We train language models on C4 tokenized with the T5 tokenizer (Raffel et al. 2020) and report results in terms of validation loss.
Models. We start from the Gal Codebase (Zhao et al. 2024a) and train L.LMMA models of three sizes: 130M, 350M, and IB. The models have widths of 768, 1024, and 2048 and depths of 12, 16, and 24. We use the 130M model to explore various ablations as shown in Section 4.1 The MLP hidden dimension of the 130M model is 4 times the width and the hidden dimension of the 350M and IB model is 8 times the width. The activation function is SiLU (Elfwing et al. 2018). The 3 architecture uses RoPE positional encodings (Su al. 2024). Attention heads are always dimension
64. For more architecture details please refer to Zhao et al. (2024a). We train in mixed precision with FP32.
Algorithms. We use the standard Pytorch implementation of Adam, and the official Gal.ore imple- mentation provided by Zhao et al. (2024a). Since Two-Sided SOAP is too memory consuming and not within comparison scope, we modify the code provided by Vyas et al. (2024) to apply One-Sided SOAP discussed in Vyas et al. (2024). We use the official Adam-mini implementation provided by Zhang al. (2024b). For MUON and NS5, we their implementation provided by Jordan et al. (2024) in their Github records. We implement our COSMOS starting from an older version of Pytorch implementation of AdamW.
Default hyperparameters. In all algorithms, we choose first order momentum 31 = 0.9 to align with and get a fair comparison with Adam baseline. We choose second order momentum 8n 0.98, which is also a widely used configuration after Liu et al. (2019) mentioned that it provides better training stability than 0.999. We set smoothing term e 1e-8 align with the standard hyperparameter choice. We use the linear learning rate schedule to decay the learning rate to O. To align with Zhao et al. (2024a), we set the warmup ratio be 10% and weight decay be
Token counts. For all of our runs we use sequence length of 1024. For the 130M model, we set the batch size be 960, and for the 350M and 1B models, we set the batch size to 2000. We train the 130M and 350M models for 5k steps and train the 1B model for 13k steps. Thus the number training tokens for the 130M mode se 5B, which is beyond the "chinchilla optimal" number of tokens. The numbers of training tokens for the 350M model and IB model are 10B and 26B respectively, which follow the chinchilla optimal" number tokens.
A.1 LEARNING RATE TUNING
To avoid unfair comparisons caused by excessive hyperparameter tuning, for all algorithms we set the learning rate as the only tunable hyperparameter in all the main results in Section The rank for COSMOS for all main results is fixed 64.
A.1.1 TUNING ON 130M MODEL
For Adam, we tune the learning rate on (2.5c-4, 5c-4, 1e-3, 2c-3, 4e-3, 8c-3). In our experiments, 2e-3 is the optimal learning rate and 8e-3 diverges. Then for SOAP, we also tune the learning rate on {5e-4, le-3, 2e-3, 4e-3). : For Adam-mini, we just use the optimal learning rate of Adam, which is also 2e-3.
For GaLore, We follow the setting in Zhao et al. (2024a), set rank=256, and scale factor a 0.25. According to Zhao et al. (024a)), the learning rate of Galore should be larger than Adam's. They mentioned that Galore not sensitive to hyperparameter and they use the same learning rate le-2 for all size of models after tuning, simply tune galore in range near le-2, which is {5c-3, le-2, 2e-2, 2c- The projection update frequency is 200 for 20k training steps, thus we decrease it to 50 for our 5k training steps.
For the implementation of MUON and COSMOS, the embedding and output layer will use Adam while other parts will use MUON/COSMOS algorithm. To avoid multiple hyperparameter tuning, we fix the learning rate for embedding and output layer to 2e-3, which is the optimal learning rate for Adam, and only tune the learning rate of hidden layers, whose optimizer is MUON/COSMOS.
14

---

Published as a conference paper at ICLR 2026
To be more specific, we tune the learning rate of hidden layers {le-4, 2c-4, 5e-4, 1e-3, 2--3). is worth noting that (Liu et al., 2025) suggests the optimal learning rate for MUON should be 0.2-0,4 times that of the Adam learning rate used for the embedding layer (2e-3 in our setting), which exactly falls within the range we searched.
For COSMOS,
mentioned before, avoid tuning 7 we simply y be the ratio of the learning rate of hidden layers to the learning rate of the embedding layer (which is fixed at 2e-3). We find that this trick can provide satisfactory result without extra tuning on y. Please note that we find in many extra experiments that this trick isn't 't the optimal choice for 7, Tuning 7 may output a better result.
A.1.2 TUNING ON 350M MODEL
For Adam and SOAP, we tune the learning rate on {2.5e-4, 5e-4, 1e-3, 2e-3, 4e-3, 8e-3), which is same as the range in Section A.I.I For Galore, we set the rank to be 384, projection update frequency to be 50, and scale factor a = 0.25. Then we tune the learning rate of GaLore on (5c-3, le-2, 2c-2, 4c-2]. which is also same what we do in Section
For the implementation of MUON and COSMOS, we still fix the learning rate for embedding and output layer to be 2e-3 and only tune the learning rate of MUON/COSMOS for hidden layers. For MUON and COSMOS, we tune the learning rate on { (1e-4, 2e-4, 5e-4, 1e-3, 20-3} to align with our setting in Section A.I.I Also for COSMOS, we still set y to be the ratio of the learning rate of hidden layers to the learning rate of the embedding layer (which is fixed at 2e-3).
1.1.3 TUNING ON 1B MODEL
We do not have enough resources to tune hyperparameters carefully on the IB model. For Adam, we first try learning rate 1) = 2e-3, but an extremely large loss spike occurred in the early stage. Then we decrease to 1e-3 and get the baseline result. For MUON and COSMOS, we still fix the learning rate for embedding and output layer to be 2e-3 and tune their learning rate on (2e-4, 5e-4). For COSMOS, we still set y to be the ratio of the learning rate of hidden layers to the learning rate of the embedding layer.
A.1.4 DISCUSSION ON LEARNING RATE TUNING
We are discussing the learning rate used by MUON in their reproducible logs here to demonstrate that our learning rate falls within reasonable range.
There are two versions of Muon implemented in the reproducible logs of modded nanogpt. In the early versions, the algorithm they used was consistent with what we described in Algorithm 3 This algorithm has been used on both 124M and 1.5B GPT models and achieved SOTA performance.
In this version, they used Adam's 's baseline learning (3.6e-3) rate as the learning rate for the embed- ding and output layers on a 124M model, and used 3.6e-4 as the cleaning rate for Muon. In our experiment, since the Adam baseline learning rate we obtained was 2e-3, which is a little different with 3. 3.6e-3, we also use this learning rate as the learning rate for the embedding and output layers. We avoid adjusting the learning rates of the embedding and output layers, as this would result in the tuning of both learning rates for two parts of parameters. Generally speaking, this would yield better results than adjusting only one learning rate, but this effect not fair compared to the Adam algorithm with only one learning rate. For Muon's 's learning rate in our experiments, our traversal set { [le-4, 2e-4, 5e-4, 1e-3, 2c-3] is also relatively close to their 3.6e-4.
In the later reproducible logs, they made a simple modification to Muon, but did not mention whether this would improve the effect. This modification to change the line 7 in Algorithm to be W4+1 < W. nB.
Considering that n for different matrix parameters in the same LLM are the same (e.g., 768 in the 124M GPT), this method is just simple rescale. However, due to the current scale being reduced by v times, this method generally requires a larger learning rate. For example, in their subsequent logs, they used 0.02 This learning rate is nearly equivalent to using 0.02/ /76 x 7.2e in the first version, which is also close our traversal set ( (1e-4, 2e-4, 5e-4, 1e-3, 2e-3).
15

---

Published as a conference paper at ICLR 2026
Since MUON did not specify which version would be more advantageous, used the first version in our experiments, as it provided more reproducible logs. At the time we began experimenting with MUON, the second version first version, which was also adopted in the subsequent work on MUON scaling (Liu et al. 2025
yet available. For consistency, therefore continued with the
A.2 ABLATION OF r AND 7 ON 130M MODEL
For the ablation experiments on COSMOS Section 4.1 we tune discount factor y and rank together to show COSMOS isn't very sensitive to hyperparameters. We fix the learning rate for em- bedding and output layers to be 2e -3, fix the learning rate for COSMOS to be 5e-4, and sweep over the cross product of r E (32, 64, 128) and E 0.25, 0.5, 1}. With all these hyperparameters COSMOS outputs comparable results to MUON.
A.3 ABLATION OF NORMALIZATION
As mentioned in the normalization paragraph in Section 4.1 performance of COSMOS than MUON is simply because the normalization function NORM, we modify the normalization method of MUON 350M models. We still tune the learning rate of MUON NORM on {Se-3, le-2, 2e-2, 4e-2]. and present their best performance.
exclude the possibility that the better
NORM and rerun the experiments for 130M and
A.4 PROFILING EXPERIMENTS
We do the profiling experiments the 1B model. We set the sequence length 1024, which aligns with our previous settings. We set batch size 10 and accumulation steps 25. Then we record the maximum GPU memory usage and time usage this setting by using Pytorch API during the entire forward-backward backward and optimizer update process.
B EXPERIMENTS DETAILS ON MODDED-NANOGPT
As discussed in Section 4.3 we directly use Muon's reproducible logs modded NanoGPT. In the setting of GPT-2 Small (1 14M), they set the learning rate for embedding layer (optimized with Adam) to be 3.6e-3, and the learning rate for hidden layers (optimized with Muon) to be 3.6e-4. Also they use Warmup-Stable Saable-eccyy (WSD) schedule instead of Cosine schedule. Their batch size is 512, sequence length is 1024 and number of iterations is 6200. We also plot the wall-clock time comparison Figure
In the training of 124M model, followed their original setting for Muon and only additionally searched 31 within (0.9, 0.95]. For COSMOS, we adopted the same setting as Muon, also searching 31 in {0.9, 0.95). For 82 and in COSMOS, we them 0.95 and 0.2 without additional search.
In the training of GPT-2 Medium (350M), they used a very uncommon setting: the learning rate for embedding layer is 0.3, which very large. But the learning rate for output layer is still 3e-3. They also reduced the momentum for the embedding layer and output layer to 0.8 which is also an uncommon choice.
To demonstrate the generality of COSMOS across various settings, our experiments on the 350M model completely follow their setting. We only additionally search 3n within {0.9.0.95} for both Muon and COSMOS. However, since this setting is indeed uncommon and subsequent work on Muon scaling (Shah al., 20255] [iu et al., 2025) did not follow it, we did not adopt this setting in our other experiments (LLaMA on C4 and GPT2 on Wikitext103).
C TWO-SIDED COSMOS
For simplicity, we only consider the one-side version of COSMOS in this paper. Like SOAP, COS- MOS can be further generalized two-sided version. Similar to two-sided SOAP in Vyas et al. 2024), provide two-sided variant of COSMOS in Algorithm
16

---

Published as a conference paper at ICLR 2026
Muon
3.40
Cosmos os
Validation
3.30
3.29
2000
2500
3000
3500
4000
4500
5000
Wall Clock time (s)
Figure 6: The validation loss vs. wall clock time plot for 350M NanoGPT trained on FineWeb.
Algorithm 4 Two-sided version of COSMOS for n layer. Per layer, maintain six matrices: U E Rnxr, O E R R,R E Rx V E R and M E p2xx.
input Learning rate 7/ combination weight 7, projection rank < momentum parameters (31, perturbation parameter e. For simplicity, we omit the initialization.
1: for t = O.... do
2: Sample batch Mr
3:
G. VwØM, (W;)
4:
Mr + 81M + (1 (1
5:
U. 2R(32U)-181-111 (1-
6:
Oz t (1
0-1)
7:
S. (6
32] G.
8:
RG + 0, 01(000-1100000
9:
V + (1-)) G())) (OT
O, M.,(/111-99) 39)
A O,
((Vi (v(V++ +c)((1 35)
U.
||M- O,M,U,, li ))
M, M2 O.M.U.TT
11: B. + NORM NS5
12:
G. A t7 BB m
13: W411 < W. - NORM(G.)
14: end for
D ADDITIONAL EXPERIMENTS
This section provides supplementary experiments not presented in the main text.
D. EXPERIMENTS FOR NORMALIZATION ABLATION
We conduct additional experiments LLaMA-130M and LLaMA-350M to validate that normal- ization is not the main reason for COSMOS outperforming MUON. As shown Figure normal- ization does not make much difference MUON, while COSMOS consistently outperforms both of them.
D.2 EXPERIMENTS ON WIKITEXT
This section discuss the details of the experiments on WikiText (Merity et al., 2016) and GPT- 2 (Radford al. 2019). To be more specific, train GPT2-small (25M) and GPT2-medinn (355M) on the Wikitext-103 dataset. We discard learnable position embeddings and use RoPE (Su al. 2024 replacement.
17

---

Published as a conference paper at ICLR 2026
3.25
Muon
Muon
noon+norm
Muon+
3.2
3.20
Cosmos
Cosmos
Validation Loss
3.1
Validation
3.0
3.10
2.9
3.05
3.00
2000
3000
4000
5000
1250
2250
3250
4250
Update Steps
Update Steps
(a) Comparison of COSMOS, MUON and MUON with normalization for LLaMA-130M on C4.
(b) Comparison of COSMOS, MUON and MUON with normalization LLaMA-350N C4.
Figure 7: Comparison of COSMOS, MUON, and MUON with normalization LLaMA-130M and LLaMA-350M for C4.
For GPT2-small, we tune the learning rate of Adam (5e-4, le-3, 2e-3, 4e-3, 8e-3), and find 4e-3 is the optimal learning rate for Adam. Then in MUON/COSMOS, we use learning rate 4e-3 for the embedding layer and 4e-4 for MUON/COSMOS.
Similarly, for GPT2-medium, we tune the learning rate of Adam on (5e-4, le-3, 2e-3, 4e-3, 8e-3), }. and find 2e-3 is the optimal learning rate for Adam. Then MUONICOSMOS, we use learning rate 2e-3 for the embedding layer and 5e-4 for MUON/COSMOS.
For COSMOS, y still set to be the ratio of the hidden layer learning rate the embedding layer learning rate in both models.
We the sequence length be 1024, and the batch size is also 1024. We train both models for 5k steps, which means the models are trained on 5B tokens. For such many training tokens on Wikitext- 103, overfitting will occur and validation loss will start to increase after training for certain number of steps. Therefore, we use the training loss as the metric for comparison.
The results for GPT2-small and GPT2-medium are provided Figures 5a and 5b respectively. We observe that COSMOS consistently outperforms MUON, showing that does not overfit any particular model or dataset.
D.3 SMALLER LEARNING RATE FOR MUON/COSMOS ON LLAMA-IB
As discussed in section A.1.3. we tune the learning rate for MUON/COSMOS on (2e-4, 5e-4). We find 5e-4 is better and present its corresponding results in the main text. Here present the result for 2e-4 in Figure g where COSMOS still outperforms MUON.
D.4 WALL-CLOCK TIME PLOT FOR LLAMA-IB
Based on the throughput we calculate in Table we rescale the X-axis of Figure 2h to be wall-clock time and present the result Figure
D.5 EXPERIMENTS ON SHORTER SEQUENCES
To validate the correctness our implementation, we compare GaLore with COSMOS and Adam on 256 sequence length adopted in /hho et al. (002a)). As shown in Figure TO Gal and Adam
18

---

Published as a conference paper at ICLR 2026
2.85
Muon
2.80
Cosmos
Validation Loss
2.75
2.70
2.65
2.60
2.55
3000
5000
7000
9000
11000
13000
Update Steps
Figure 8: LLaMA-1B trained on C4 dataset with learning rate 2e-4 for MUONICOSMOS. COS- MOS still outperforms MUON.
2.75
Adam
Muon
Cosmos
2.65
Validation
2.60
2.55
2.50
Wall clock time (Hours) on 2AA100
Figure 9: Wall-Clock time plot for our training on LLaMA-1B
are more comparable in this shorter sequence setting, suggesting our implementation is correct and degradation GaLore shown Figure and Ta mainly due long sequence length.
Adam
Galore
Validation Loss
Cosmos
2050 cole séco sodee 12000 sodco seloo and00 zolco Update Steps
Figure 10: Comparison COSMOS, Adam, and GaLore on 256 sequence length. The performance of GaLore on shorter sequences does not deteriorate for long sequences, validating the correctness of our implementation.
19

---

Published as a conference paper at ICLR 2026
E LLM USAGE
In preparing this paper, large language models (LLMs) such as ChatGPT were used only for light editing purposes, including minor grammar checking and sentence polishing. No part of the search ideation, methodology design, experimental execution, or analysis was conducted with the assistance of LLMs.
20


Reviews:
[{"review_id": "V0PwHeN0xi", "reviewer": "Reviewer_BuFv", "summary": "COSMOS presents an innovative hybrid optimization approach that cleverly combines SOAP and MUON optimizers through eigenspace decomposition. The authors tackle the important challenge of reducing memory consumption in LLM training while maintaining optimization quality. By applying SOAP to the leading eigensubspace and MUON to the remaining space, the method achieves notable efficiency gains with reported memory savings of 82% and training time reduction of 91%.", "questions": "-- Could you provide theoretical analysis showing when COSMOS converges faster than pure SOAP/MUON?\n\n-- Have you explored adaptive rank selection based on the eigenvalue spectrum?\n\n-- Could you test on at least one non-transformer architecture to demonstrate generalizability?\n\n-- Is there a principled way to choose γ without the η/η0 heuristic?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "-- Addresses a critical problem: Memory efficiency in LLM training is a major bottleneck, and this work makes genuine progress\n\n-- Creative solution: The eigenspace decomposition approach shows good intuition about balancing computational cost and optimization quality\n\n-- Promising empirical results: The reported efficiency gains are substantial and could have real practical impact\n\n-- Clear presentation of core ideas: The authors effectively communicate their hybrid strategy", "weaknesses": "-- Theoretical gaps: The paper would greatly benefit from convergence analysis and theoretical justification for the eigenspace splitting strategy. Adding even basic convergence guarantees would strengthen the contribution significantly\n\n-- Limited scope validation: Expanding beyond transformer architectures to CNNs or other models would demonstrate broader applicability\n\n-- Hyperparameter sensitivity: The rank r and weight γ selection needs more principled guidelines - perhaps adaptive schemes could be explored\n\n-- Missing key ablations: Comparing different eigenspace decomposition strategies or investigating the optimal split ratio would provide valuable insights", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "VAMkUM9APl", "reviewer": "Reviewer_NctS", "summary": "See below", "questions": "The paper proposed a new pre-training optimizer called COSMOS, which has similar performance as SOA  but has cheaper memory requirement than SOAP. The key idea is to compress the SOAP’s second-moment matrix using its dominant eigenspace and track it in a low-dimensional subspace.\n\nThe idea is interesting. The reduction in the complexity of SOAP is potentially useful.   The presentation is clear. \n\nI have concerns for several experimental settings adopted in the current script, which I think might not fully reflect the performance of the proposed method. I suggest the following setup changes and ablations. \n\n1. **Non-zero weight decay.** The authors use weight dacay = 0. What if we use non-zero weight decay such as weight decay = 0.1? The choice of weight decay can be crucial to the optimizer behaviors, as pointed out by recent benchmark [1].\n\n2. **More training tokens.** For Fig 1: 5k step (5B tokens) is roughly 1x chinchilla. How would the methods perform under an overtrained setting, such as 4x or 8x chinchilla in the recent work [1]? Will the advantage of COSMOS degenerate, or will it be surpassed by Muon or other methods? It would be helpful to see a new version of Fig. 1 under 4x or 8x chinchilla setting. \n\n3. **More ablation on batchsize.** Matrix-based methods usually can afford a larger batchsize than Adam-type methods.  I think the ablation on batch size is insufficient in the current script. It is possible that the optimizer ranking becomes quite different under different batchsize settings. It is also possible that the proposed method performs even better in a larger batch setting. \n\n4. **Higher-quality datasets.**  C4 is popular,  but the data quality is relatively poor and noisy from the current standard. I think the ablation on the more SOTA dataset is needed. I understand it might be too expensive to re-run all the experiments on Fineweb, but I think at least some ablation on the small-scaled experiments is needed.  For the smallest experiments in the paper (125M model), can the authors provide additional ablation on the dataset by changing the dataset to Fineweb and re-evaluate all the methods?  Will the data quality affect the optimizer rankings? \n\nI am happy to raise my score if my concerns are addressed. \n\n[1] Fantastic Pretraining Optimizers and Where to Find Them", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "See below", "weaknesses": "See below", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "SBqTxDEkN6", "reviewer": "Reviewer_FdBR", "summary": "This paper proposes COSMOS, a hybrid optimizer that combines the strengths of SOAP and MUON for large language model (LLM) training.  \nThe core idea is to decompose the gradient matrix into leading eigensubspace and residual subspace:  \n- A SOAP-like update is applied to the leading subspace to capture major optimization dynamics.  \n- A MUON-like update is used for the residual subspace to reduce memory and computation costs.  \nThis approach greatly reduces memory usage (about 19% of SOAP) without noticeable performance loss.  \n\nExperiments on LLaMA, GPT-2, and NanoGPT show that COSMOS outperforms MUON in token and memory efficiency while achieving performance comparable to SOAP.", "questions": "How is the rank determined? Is it based on an energy ratio or an information-based criterion?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Innovative idea: A hybrid low-rank optimization framework that balances performance and resource cost.\n\nComprehensive experiments: Covers multiple models, datasets, and baselines.\n\nDetailed analysis of memory and efficiency: Provides quantitative evidence of significant savings.\n\nStrong robustness: Insensitive to rank and hyperparameter variations.", "weaknesses": "Insufficient methodological details:\nLacks analysis on rank selection and the dynamics of subspace changes during training.\n\nWeak theoretical explanation:\nUnclear whether subspace variation causes instability in the optimization direction.\n\nLimited experimental analysis:\nLacks mechanistic or spectral analysis to explain why COSMOS outperforms SOAP.\nThe fundamental advantage of COSMOS over SOAP and MUON remains unclear, and the observed performance degradation with higher ranks is only heuristically explained.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "k3DgiP0Fyc", "reviewer": "Reviewer_isQ9", "summary": "The paper proposes a new optimizer named COSMOS, which applies the update of SOAP on leading eigensubspace of parameters and applies Muon for the rest. The motivation for the method is to reduce the memory cost of storing second moment of SOAP. Numerical results show that COSMOS is able to achieve comparable convergence as SOAP in the task of pre-training LLMs under Llama and GPT-2 models, while costs less memory.", "questions": "* Could the authors provide more comments on why setting $\\gamma=\\frac{\\eta}{\\eta_0}$? The $\\gamma$ determines the weight of Muon's update; it is unclear why this is related to the learning rate in embedding/language modeling head layers.\n\n* Given the fact that COSMOS is faster than Muon, it would be interesting to see the effect of applying SOAP on different range of eigenspaces rather than the top eigenspace, which helps understanding the algorithm more comprehensively.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "* The proposed method avoids the storage of second moment in SOAP and is memory-efficient.\n* The method achieves comparable convergence speed as SOAP in pre-training Llama/GPT-2 models of different scales, and is shown to be robust to newly introduced hyper-parameters, e.g. $\\gamma$.", "weaknesses": "* The main motivation for the design of COSMOS is not well presented. It is unclear why replacing the top eigenspace's update in Muon with SOAP will accelerate the convergence. It has been shown that SOAP exhibits comparable in-iteration convergence as Muon [1].\n\n* The acceleration of COSMOS over Muon is marginal. Given COSMOS' higher memory cost and training FLOPs, the practical benefit is a bit marginal compared to Muon.\n\n[1] Muon: An optimizer for hidden layers in neural networks", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "dmiMSG3tx7", "reviewer": "Reviewer_vF7y", "summary": "The paper introduces an efficient optimizer for pre-training LLMs by combining two previously successful optimization methods. Specifically, it applies SOAP to the leading eigenspace of the gradient matrix and Muon to its orthogonal complement. This hybrid approach is claimed to yield a more memory-efficient and faster optimizer than SOAP, while also outperforming Muon. The authors validate their method through LLM training experiments on standard benchmarks, demonstrating improvements over baselines.", "questions": "- (1) and (2) in the introduction are somewhat vague. In (2), you mention a limitation related to storing adaptive learning rates —what do you mean here? It would be more precise to state that the optimizer stores both the first and second moments, which increases memory consumption.\n- On the line 167 you say G-S procedure. Do you implement your own QR with G-S? Presumable, you use torch.qr here which doesn't use G-S. \n- Line 220, regarding Muon's scaling, I think the right scaling for learning rate transfer is $\\sqrt{\\frac{m}{n}}$. Am I missing something?\n- In table 1 would be good to see COSMOS memory usage dependence on $r$ as well.\n- In Figure 1 and subsequent experiments, why do you observe SOAP to be worse then COSMOS. Isn't it counterintuitive?\n- I suspect forward/backward takes most of the time in Table5. Can you separate out the optimizer step time in that plot? That will give a fair comparison. \n- Can you provide wall-clock plots for nano-gpt experiments?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The idea of applying SOAP to the leading eigenspace and Muon to its orthogonal complement is useful and promising. \n\n - Leverages the strength of a stepwise effective yet computationally expensive optimizer by restricting its use to a lower-dimensional subspace, and hence improving memory and computational efficiency. \n\n- This approach could inspire further exploration of hybrid methods that combine different optimizers using their principle. \n\n- The paper presents pretraining experiments across standard model sizes and datasets, demonstrating improvements of the proposed method.", "weaknesses": "- The main concern is that the paper does not provide sufficient wall-clock time experiments, making the claimed advantage over Muon questionable. The only such experiment involves a 1B LLaMA model. However, due to the large gradient accumulation, I think the forward and backward passes dominate the computation time, making the optimizer’s step time negligible. As a result, the reported plot may effectively reflect only a stepwise comparison rather than a true wall-clock evaluation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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:

Decision: Accept (Poster)