Paper deep dive
Clean: Second-order LLM Training at Linear Memory Cost via Nyström Sketching
Beheshteh T. Rakhshan, Sahar Rajabi, Maziar Sargordi Shikai Fang, Guillaume Rabusseau, Sirisha Rambhatla
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 91%
Last extracted: 10/8/2026, 6:08:06 AM
Summary
The paper introduces Clean, a memory-efficient second-order optimizer for Large Language Models (LLMs) that approximates the full-curvature SOAP optimizer using randomized Nyström sketching. Clean reduces memory complexity from quadratic to linear in model dimensions by maintaining low-rank approximations of preconditioners and applying block-wise Adam updates. A quantized variant, Q-Clean, further reduces memory usage. Clean achieves performance competitive with SOAP and AdamW but with significantly lower memory footprint, enabling the training of 13B-parameter models on a single 80GB GPU.
Entities (8)
Relation Signals (6)
Q-Clean → isvariantof → Clean
confidence 95% · We further propose Q-Clean, a low-precision variant that aggressively compresses optimizer states.
Clean → usesmethod → Nyström Sketching
confidence 95% · Clean leverages the randomized Nystrom method to accurately approximate the left and right preconditioners in SOAP
Clean → enablestrainingon → H100 GPU
confidence 90% · enable the pre-training of a 13B-parameter model on a single 80GB GPU
Clean → isvariantof → SOAP
confidence 90% · Clean is the first second-order SOAP-like optimizer whose memory grows only linearly
Clean → outperforms → AdamW
confidence 85% · reaching AdamW's final performance 26% faster in wall-clock time
Q-Clean → reducesmemoryvs → Muon
confidence 85% · Q-Clean reduces optimizer memory consumption by over 50% compared to Muon
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:Training large language models (LLMs) entails a fundamental trade-off: memory-efficient optimizers such as Adam discard cross-parameter curvature, whereas full-curvature methods such as SOAP can accelerate convergence at prohibitive memory costs. We introduce Clean, a memory-efficient and full-curvature optimizer designed to resolve this bottleneck. Clean leverages the randomized Nystrom method to accurately approximate the left and right preconditioners in SOAP, and to reduce the optimizer's memory complexity from quadratic to linear in terms of model dimensions. We subsequently reintegrate the off-subspace components to capture curvature information beyond the low-rank approximation, preserving rich curvature at minimal memory cost. We further propose Q-Clean, a low-precision variant that aggressively compresses optimizer states. Q-Clean reduces optimizer memory consumption by \textbf{over 50\%} compared to Muon when pre-training a LLaMA-1.3B architecture, all while maintaining strong and competitive predictive performance. Notably, Clean operates with a smaller optimizer-state footprint than standard AdamW while reaching AdamW's final performance \textbf{26\% faster} in wall-clock time. Furthermore, our methods uniquely enable the pre-training of a 13B-parameter model on a single 80GB GPU, providing a scalable, efficient, and accessible approach to large-scale model optimization.
Tags
Links
- Source: https://arxiv.org/abs/2610.04204v2
- Canonical: https://arxiv.org/abs/2610.04204v2
Trouble viewing inline? Open PDF directly →
Full Text
71,773 characters extracted from source content.
Expand or collapse full text
Clean: Second-order LLM Training at Linear Memory Cost via Nyström Sketching Beheshteh T. Rakhshan∗*,1 Affiliation: Shikai Fang4, Guillaume Rabusseau1,5, Sirisha Rambhatla2 Affiliation: 1Mila & DIRO, Université de Montréal Sahar Rajabi∗*,2 Affiliation: 2Critical ML, University of Waterloo Affiliation: 3Independent Researcher Maziar Sargordi3 Affiliation: 4Zhejiang University Affiliation: 5Canada CIFAR AI Chair Abstract Training large language models (LLMs) entails a fundamental trade-off: memory-efficient optimizers such as Adam discard cross-parameter curvature, whereas full-curvature methods such as SOAP can accelerate convergence at prohibitive memory costs. We introduce Clean, a memory-efficient and full-curvature optimizer designed to resolve this bottleneck. Clean leverages the randomized Nyström method to accurately approximate the left and right preconditioners in SOAP, and to reduce the optimizer’s memory complexity from quadratic to linear in terms of model dimensions. We subsequently reintegrate the off-subspace components to capture curvature information beyond the low-rank approximation, preserving rich curvature at minimal memory cost. We further propose Q-Clean, a low-precision variant that aggressively compresses optimizer states. Q-Clean reduces optimizer memory consumption by over 50% compared to Muon when pre-training a LLaMA-1.3B architecture, all while maintaining strong and competitive predictive performance. Notably, Clean operates with a smaller optimizer-state footprint than standard AdamW while reaching AdamW’s final performance 26% faster in wall-clock time. Furthermore, our methods uniquely enable the pre-training of a 13B-parameter model on a single 80GB GPU, providing a scalable, efficient, and accessible approach to large-scale model optimization. $*$$*$footnotetext: Equal contribution. 1 Introduction Figure 1: Measured memory consumption for pre-training a LLaMA-13B model on a single 80 GB H100 GPU (maximum sequence length of 1024, batch size of 1). All non-quantized matrices are stored in bfloat16. Activation + Other combines activation memory with an estimated optimizer workspace. Methods marked ∗ ran out of memory, so their optimizer state is estimated. Training large language models (LLMs) forces a trade-off between memory and curvature. Adaptive optimizers store moment statistics that match or exceed the size of the model itself, and second-order methods such as Shampoo (Gupta et al., 2018) and SOAP (Vyas et al., 2024) store even more. The remedies that reduce this footprint typically discard curvature information, producing poorly conditioned updates and slowing convergence (Gupta et al., 2018). Balancing memory efficiency with curvature fidelity has therefore emerged as a central challenge for efficient LLM training. A range of optimizers have been proposed, each making a different trade-off between memory and curvature. First–order methods such as Adam (Kingma and Ba, 2014) and Adafactor (Shazeer and Stern, 2018) are lightweight optimizers that store only diagonal statistics. Their low memory cost makes them practical for large models. Full–matrix second–order methods such as Shampoo (Gupta et al., 2018) and KFAC (Martens and Grosse, 2015) explicitly capture curvature and can accelerate training, but their quadratic state and computation make them prohibitive at LLMs scale. More recently, SOAP (Vyas et al., 2024) achieved highly stable and full-rank preconditioning by applying Adam to the gradient projected onto the eigenbasis of Shampoo’s preconditioner; however, it still relies on explicitly accumulating full-rank covariance matrices. Several efficient optimizers have been proposed, but all of them have at least one optimizer memory state that grows quadratically with respect to model dimensions (i.e., (mn)O(mn)) (Romashov et al., 2026; Jordan et al., 2024; Gong et al., 2026; Liu et al., 2025), making them difficult to scale. To resolve this, we introduce Clean, a curvature-aware optimizer that achieves SOAP-level performance while its memory states scale linearly with respect to model dimensions. Clean replaces explicit gradient covariance accumulations with randomized Nyström approximations, extracting highly accurate, compact estimations for the left and right Shampoo preconditioners. It then splits the update into four orthogonal blocks, Main, A, B, and C, and applies full Adam to Main. Rather than discarding gradients that fall outside the low-rank core, it applies pooled Adam-like updates to the three remaining blocks. Finally, every block is projected back to the parameter space, and Clean recombines all four to produce a rich and full-rank update, as shown in Fig. 2. We also introduce Q-Clean, a quantized variant that aggressively compresses the stored optimizer states. To the best of our knowledge, Clean is the first second-order SOAP-like optimizer whose memory grows only linearly with layer dimensions, requiring ((m+n)r)O((m+n)r) instead of (mn)O(mn) for a layer of size m×nm× n. Extensive pre-training evaluations on LLaMA-based architectures demonstrate that our approach successfully retains the benefits of full-rank second-order preconditioning at a fraction of the optimizer memory cost. Clean delivers validation perplexities highly competitive with SOAP, yet operates with a smaller memory footprint than AdamW, while showing significant speedup in comparison to Adam’s convergence rate. Fig. 1 shows that across all the baselines, Clean and Q-Clean are the only optimizers that allow training a 13B model on an H100 GPU with 80GB of memory. Figure 2: Clean. (I) GtG_t is projected via the estimated left and right preconditioners to yield GtpreG_t^pre, which is then partitioned into four blocks Main, A, B, and C. (I) Adam-style update is applied to each block, maintaining either full or pooled moments depending on the block’s dimensions (only the shaded buffers are stored). (I) The Adam output from each block is projected back into the parameter space and summed to form a full-rank update. In summary, Clean is the first optimizer to perform SOAP-style preconditioning in the eigenbasis of randomized Nyström approximations of both Kronecker factors. The Nyström approximations allow the covariance factors to be updated without ever forming them explicitly, reducing their memory cost from quadratic to linear in the layer dimensions for a fixed rank r. Even when the rank increases with the model dimensions, Clean reduces the per-parameter cost as the model scales. Finally, we present Q-Clean, a low-precision extension that quantizes optimizer states to achieve competitive performance while further reducing optimizer memory requirements. 2 Background and Preliminaries Most deep learning optimizers can be viewed as preconditioned gradient methods. Consider a weight matrix Wt∈ℝm×nW_t ^m× n with vectorization wt=vec(Wt)∈ℝmnw_t=vec(W_t) ^mn, loss function ϕ(⋅)φ(·) and stochastic gradient gt=∇ϕt(wt)∈ℝmng_t=∇ _t(w_t) ^mn. A generic preconditioned update takes the form wt+1=wt−ηPt−1gt,w_t+1=w_t-η\,P_t^-1g_t, (1) where η>0η>0 is the learning rate and Pt∈ℝmn×mnP_t ^mn× mn is a positive semi-definite (PSD) preconditioner that encodes curvature information. Different optimizers correspond to different choices of PtP_t, balancing the accuracy of the curvature with memory and computational costs. We first review several widely used optimizers, and then introduce the technical tools on which our method builds. 2.1 Optimizers as Preconditioners Eqn. (1) provides a common lens for comparing optimizers: each corresponds to a particular choice of PtP_t, which dictates the balance between retained curvature information and memory requirements. Adaptive optimization as preconditioning. Full-matrix AdaGrad (Duchi et al., 2011) sets Pt=∑s=1tgsgs⊤P_t= _s=1^tg_sg_s , yielding well-conditioned updates but requiring (m2n2)O(m^2n^2) memory, which is infeasible for modern models with billions of parameters. Common approaches simplify PtP_t to make optimization practical. Diagonal preconditioners, used in optimizers such as Adam (Kingma and Ba, 2014), restrict PtP_t to a diagonal matrix, capturing individual parameter scales but ignoring cross-parameter correlations, decreasing the optimizer memory complexity to (mn)O(mn). Balanced preconditioning for matrix layers. A practical surrogate to approximate Pt∈ℝmn×mnP_t ^mn× mn is maintaining left and right covariance accumulators: Lt=∑s=1tGsGs⊤,Rt=∑s=1tGs⊤Gs,L_t= _s=1^tG_sG_s , R_t= _s=1^tG_s G_s, (2) where GtG_t is the gradient of W at step t. The preconditioned gradient is then G~t=Lt−1/4GtRt−1/4, G_t=L_t^-1/4\,G_t\,R_t^-1/4, (3) which captures the row and column curvature without ever forming the full preconditioner (Gupta et al., 2018). Concretely, PtP_t is approximated by the Kronecker product of two smaller factors, LtL_t and RtR_t, where the left preconditioner Lt∈ℝm×mL_t ^m× m acts on the rows of GtG_t, while Rt∈ℝn×nR_t ^n× n acts on its columns. Shampoo (Gupta et al., 2018), a popular second-order optimizer, takes exactly this route and has become the standard Kronecker-factored approximation to full-matrix AdaGrad (Duchi et al., 2011), which computes and stores per-layer LtL_t and RtR_t, utilizing fractional powers Lt−α,Rt−αL_t^-α,R_t^-α as left and right preconditioners. The original α=1/4α=1/4 matches AdaGrad in theory, but Anil et al. (2020) found that α=1/2α=1/2 often improves stability. However, fully storing LtL_t and RtR_t still incurs quadratic memory in the layer width, which is prohibitive for wide layers. SOAP and eigenspace optimization. Rather than explicitly computing the costly inverse fractional powers of LtL_t and RtR_t at every step, SOAP (Vyas et al., 2024) periodically computes their eigen decompositions. It then projects the gradients into this slowly changing principal curvature eigenspace, applies a coordinate-wise Adam update, and projects the step back to the original parameter space. Ultimately, this hybrid approach unifies the structural curvature benefits of second-order optimization with the low per-step cost of diagonal methods, empirically reducing both the number of iterations and the wall-clock time compared to standard Shampoo and AdamW. However, SOAP’s memory footprint remains (m2+n2)O(m^2+n^2), the same as Shampoo’s, since it must still store the full m×m× m and n×n× n Kronecker factors, LtL_t and RtR_t at each step. The SOAP algorithm is provided in Appendix C. 2.2 Randomized Nyström preconditioner The randomized Nyström approximation (Gittens and Mahoney, 2016) provides a framework for efficiently representing symmetric PSD matrices using low-rank approximations. Given a symmetric PSD matrix A∈ℝn×nA ^n× n and a random (Gaussian) matrix Ω∈ℝn×r ^n× r where r is a rank parameter, the randomized Nyström approximation is given by: Y=AΩ,B=Ω⊤AΩ,A~=YB†Y⊤≃A,Y=A , B= A , A=YB Y A, (4) where A~ A is a rank-r surrogate, requiring only (nr)O(nr) storage instead of (n2)O(n^2) when r≪nr n. Strong theoretical guarantees (Gittens and Mahoney, 2016) and practical efficiency of Nyström approximations, make them attractive for optimizer states. However, a naive implementation based on Eqn. (4) is numerically unstable (Frangella et al., 2023). Alg. 2 in Appendix A presents a numerically stable approach for constructing a Nyström approximation (Tropp et al., 2017). Given the eigenvalue decomposition A~=UΛU⊤ A=U U returned by Alg. 2, where U∈ℝn×rU ^n× r has orthonormal columns and Λ=diag(λ1,…,λr) =diag( _1,…, _r) with λ1≥⋯≥λr≥0 _1≥…≥ _r≥ 0, the Nyström preconditioner is obtained by: P=(λr+μ)−1U(Λ+μI)U⊤+(I−UU⊤).P=( _r+μ)^-1\,U( +μ I)U +(I-U ). (5) In effect, P−1P^-1 scales down each of the top-r eigendirections of A by the factor (λr+μ)/(λi+μ)≤1( _r+μ)/( _i+μ)≤ 1, so that directions with larger eigenvalues are scaled down more, while acting as the identity on U⟂U . 3 Methodology In this section, we present our second-order memory-efficient optimizer, Clean, starting by approximating the left and right preconditioners in SOAP via the Nyström method. During training, the singular values of the gradients decay rapidly, concentrating most of their energy in a small subspace (Zhao et al., 2024). Since the eigenvalues of GtGt⊤G_tG_t and Gt⊤GtG_t G_t are the squared singular values of GtG_t, these covariances exhibit an even sharper spectral decay, making the SOAP factors LtL_t and RtR_t natural targets for low-rank approximation. Clean exploits this structure to reduce SOAP’s (m2+n2)O(m^2+n^2) memory to ((m+n)r)O((m+n)r), which is linear in m and n. Smoothed covariances. Similar to SOAP, Clean preconditions with exponential moving averages of the gradient covariances (Eqn. (6)). However, it never stores these covariances in full, instead, it maintains the low-rank approximations as in Eqn. (7), Lt=β2Lt−1+(1−β2)GtGt⊤,Rt=β2Rt−1+(1−β2)Gt⊤Gt.L_t= _2\,L_t-1+(1- _2)\,G_tG_t , R_t= _2\,R_t-1+(1- _2)\,G_t G_t. (6) The averaging prevents abrupt changes of basis across refreshes. For β2∈[0,1] _2∈[0,1], each update is a convex combination of PSD matrices, so LtL_t and RtR_t remain PSD and therefore admit randomized Nyström approximations (Frangella et al., 2023). Low-rank covariance factors. Clean never forms LtL_t or RtR_t explicitly. Instead, it maintains rank-r1r_1 and rank-r2r_2 Nyström approximations Lt≈Ut,LΛt,LUt,L⊤,Rt≈Ut,RΛt,RUt,R⊤, L_t≈ U_t,L _t,LU_t,L , R_t≈ U_t,R _t,RU_t,R , (7) where Ut,L∈ℝm×r1U_t,L ^m× r_1 and Ut,R∈ℝn×r2U_t,R ^n× r_2 have orthonormal columns, Λt,L∈ℝr1×r1 _t,L ^r_1× r_1 and Λt,R∈ℝr2×r2 _t,R ^r_2× r_2, are diagonal and non-negative, and r1,r2≪min(m,n)r_1,r_2 (m,n). Clean only stores thin factors Ut,L,Λt,LU_t,L, _t,L and Ut,R,Λt,RU_t,R, _t,R. To update them at step t, we replace Lt−1L_t-1 in the Eqn. (6) with its approximation and access the result only through a random sketch. Let ΩL∈ℝm×(r1+p) _L ^m×(r_1+p) and ΩR∈ℝn×(r2+p) _R ^n×(r_2+p) be Gaussian random matrices, where p is an oversampling parameter (Halko et al., 2011). The left sketch is LtΩL≈β2Ut−1,LΛt−1,L(Ut−1,L⊤ΩL)+(1−β2)Gt(Gt⊤ΩL),L_t _L≈ _2\,U_t-1,L _t-1,L (U_t-1,L _L )+(1- _2)\,G_t (G_t _L ), (8) and RtΩR_t _R is computed analogously (see Appendix A for details of Nyström approximations). Both require only thin matrix products with the stored factors and gradient, so no m×m× m or n×n× n matrix is ever formed. Applying Alg. 3 in Appendix A to these sketches yields the updated factors in Eqn. (7). Without loss of generality, we assume r1=r2=r_1=r_2=r. SOAP accumulates LtL_t and RtR_t at every step, whereas Clean accumulates them only every T steps. Moreover, it only stores the low-rank factors, at a cost of (m+n+2)r(m+n+2)r parameters, entirely eliminating the quadratic dependence of memory on the layer dimensions. As in SOAP, the factors are refreshed every T iterations and reused in between, so Lt−1L_t-1 and Rt−1R_t-1 always refer to the most recent refresh. Block-wise low-rank preconditioning. Following (Gupta et al., 2018; Morwani et al., 2024), we precondition the gradient as Gtpre=Lt−1/2GtRt−1/2G_t^pre=L_t^-1/2G_tR_t^-1/2. To expand GtpreG_t^pre with the Nyström preconditioner in Eqn. (5), for S∈L,RS∈\L,R\, we write PS=Ut,SUt,S⊤,PS⟂=I−PS,DS=(Λt,S+μI)−1/2,cS=(λS,r+μ)1/2,P_S=U_t,SU_t,S , P_S =I-P_S, D_S=( _t,S+μ I)^-1/2, c_S=( _S,r+μ)^1/2, (9) where λS,r _S,r is the smallest retained eigenvalue. This gives Gtpre G_t^pre =cLcRUt,LDLUt,L⊤GtUt,RDRUt,R⊤⏟Main+cRPL⟂GtUt,RDRUt,R⊤⏟A = c_Lc_R\,U_t,LD_L\, [rgb]0.1211,0.4648,0.707U_t,L G_tU_t,R\,D_RU_t,R _ Main\;+\; c_R\, [rgb]0.8398,0.4297,0.0781P_L G_tU_t,R\,D_RU_t,R _ A +cLUt,LDLUt,L⊤GtPR⟂⏟B+PL⟂GtPR⟂⏟C. +\; c_L\,U_t,LD_L\, [rgb]0.1719,0.5898,0.1719U_t,L G_tP_R _ B\;+\; [rgb]0.4688,0.4688,0.4688P_L G_tP_R _ C. (10) Running Adam directly on GtpreG_t^pre would require (mn)O(mn) memory for its moments. Clean instead keeps separate moment statistics for each block of Eqn. (10) according to its size. For the Main block, Adam runs on the two-sided projection Ut,L⊤GtUt,R∈ℝr×rU_t,L G_tU_t,R ^r× r. In addition, the cross components in blocks A and B lie outside one subspace but inside the other, so they are also low-dimensional: Gt,A=PL⟂GtUt,R∈ℝm×r,Gt,B=Ut,L⊤GtPR⟂∈ℝr×n.G_t, A=P_L G_tU_t,R ^m× r, G_t, B=U_t,L G_tP_R ^r× n. (11) Only the component in block C is orthogonal to both subspaces and is a full m×nm× n matrix. Rather than discarding the components outside the Main block, Clean applies Adam-style updates to each block with its own moment estimates. However, applying Adam to these blocks as in Romashov et al. (2026), requires buffers of size (r2)O(r^2), (mr)O(mr), (rn)O(rn), and (mn)O(mn) for Main, A, B, and C, respectively. The first three are low-rank, but the last is as large as running Adam directly on GtpreG_t^pre, reintroducing exactly the cost that Clean aims to remove. (a) Attention-Query Proj. (b) Attention-Key Proj. (c) Attention-Value Proj. (d) Attention-Output Proj. (e) MLP-Gate Proj. (f) MLP-Up Proj. (g) MLP-Down Proj. (h) Legend Figure 3: Fraction of the gradient norm captured in each of the four blocks. Generally, Main and A are dominant; the contributions of B and C slowly increase only towards the end of training, where the gradients become noisier. Values are averaged across all 24 instances of a layer type in a LLaMA-350M architecture. Fig. 3 illustrates the fraction of the gradient norm, averaged across all 24 instances of a layer type in a LLaMA-350M, for each of the blocks in Eqn. (10). Main dominates across nearly all layers, while C consistently contributes less. The figure also underscores the necessity of reintegrating information from blocks other than the Main, as they carry a substantial portion of the training signal, particularly in later iterations. We therefore tailor the moment state to each block, as summarized in Table 1. Table 1: Moments per block. β[⋅]E_β[·] denotes an exponential moving average with rate β∈β1,β2β∈\ _1, _2\, and memory counts stored floats. Main keeps both full moments; A and B keep full first moments, but sum their second moments over the long dimension, leaving one value per column or row; C keeps no first moment and a single scalar second moment, since per-entry states would cost (mn)O(mn). Block Adam’s Input 1st1^st Moment M 2nd2^nd Moment V Memory Main St=Ut,LGtUt,R∈ℝr×rS_t=U_t,L TG_tU_t,R ^r× r β1[St]E_ _1[S_t] β2[St⊙2]E_ _2[S_t 2] 2r22r^2 A Gt,A∈ℝm×rG_t,A ^m× r β1[Gt,A]E_ _1[G_t,A] β2[∑i(Gt,A)ij2]E_ _2 [ _i(G_t,A)_ij^2 ] mr+rmr+r B Gt,B∈ℝr×nG_t,B ^r× n β1[Gt,B]E_ _1[G_t,B] β2[∑j(Gt,B)ij2]E_ _2 [ _j(G_t,B)_ij^2 ] rn+rrn+r C Ct∈ℝm×nC_t ^m× n none, raw CtC_t β2[‖Ct‖F2]E_ _2 [\|C_t\|_F^2 ] 11 For the Main block, we apply regular Adam, maintaining entry-wise first and second moments (M and V), which each require (r2)O(r^2) memory. Mt←β1Mt−1+(1−β1)St,Vt←β2Vt−1+(1−β2)St⊙2,M_t← _1M_t-1+(1- _1)S_t, V_t← _2V_t-1+(1- _2)S_t 2, (12) where ⊙ denotes the Hadamard product. For blocks A and B, we retain the full-resolution first moments but spatially pool the second moments. Specifically, by assuming a shared scale for each column in A and each row in B, we reduce the second moment to a single scalar per column or row. Mt,A←β1Mt−1,A+(1−β1)Gt,A,vt,A←β2vt−1,A+(1−β2) 1⊤(Gt,A⊙2)Mt,B←β1Mt−1,B+(1−β1)Gt,B,vt,B←β2vt−1,B+(1−β2)(Gt,B⊙2) split&M_t, A← _1M_t-1, A+(1- _1)G_t, A, v_t, A← _2v_t-1, A+(1- _2)\,1 (G_t, A 2)\\ &M_t, B← _1M_t-1, B+(1- _1)G_t, B, v_t, B← _2v_t-1, B+(1- _2)\,(G_t, B 2)1 split (13) Finally, for block C, we omit the first moment entirely, utilizing the raw gradient instead, and employ its squared Frobenius norm as a scalar second moment. Projection-aware adaptive moments. On subspace-refresh steps, when the bases Ut,SU_t,S change (S∈L,RS∈\L,R\), the moments will get misaligned. The first moments M, MAM_ A, and MBM_ B are valid operators in the original coordinates, and they can be moved exactly into the new coordinates by: M←(Ut,L⊤Ut−1,L)M(Ut−1,R⊤Ut,R),MA←Ut,L⟂MA(Ut−1,R⊤Ut,R),MB←(Ut,L⊤Ut−1,L)MBUt,R⟂, split&M←(U_t,L U_t-1,L)M(U_t-1,R U_t,R),\\ M_ A← U_t,L &M_ A(U_t-1,R U_t,R), M_ B←(U_t,L U_t-1,L)M_ BU_t,R , split (14) However, the second moments cannot be exactly transformed. V is an exponential moving average of element-wise squared gradients, tracking only the diagonal (per-coordinate) variances. Since element-wise squaring does not commute with basis rotations, calculating the true rotated variances would require unrecorded off-diagonal cross-terms. Similar to SOAP (Vyas et al., 2024), we explicitly rotate the first moment M, and leave V in place. By relying on warm-start refreshes (see Appendix A), we constrain the overall frame drift, ensuring the stale-basis error in V remains sufficiently small. For A and B, we transport the pooled second moments on change of basis. Since each value stores the squared gradient along one direction, we use the squared overlap between the old and new bases: vA⊤←vA⊤(Ut−1,R⊤Ut,R)⊙2,vB←(Ut,L⊤Ut−1,L)⊙2vB.v_ A ← v_ A (U_t-1,R U_t,R) 2, v_ B←(U_t,L U_t-1,L) 2v_ B. (15) Finally, the scalar vCv_ C does not need to be rotated because it has no direction associated with it. Projection back and parameter update. After applying the block-wise Adam-style updates, each block is projected back to the original parameter space to construct the final weight update, as shown in Eqn. (10). Note that the momentum terms used in this step are bias-corrected, and similar to SOAP Vyas et al. (2024), we are using Adam scaling factors while projecting back, which eliminates the need for DSD_S and cSc_S factors in reconstructing GtpreG_t^pre (10). G~t=Ut,L(Mt^,Vt^)Ut,R⊤⏟BlockMain(Full Adam)+(M^t,A,v^t,A)Ut,R⊤⏟BlockA(Pooled Adam)+Ut,L(M^t,B,v^t,B)⏟BlockB(Pooled Adam)+(C^t,v^t,C)⏟BlockC(Pooled Adam), G_t= U_t,LN( M_t, V_t)U_t,R _Block~ Main~(Full Adam)+ N( M_t, A, v_t, A)U_t,R _Block~ A~(Pooled Adam)+ U_t,LN( M_t, B, v_t, B)_Block~ B~(Pooled Adam)+ N( C_t, v_t, C)_Block~ C~(Pooled Adam), (16) where (M^,V^)=M^/(V^+ϵ)N( M, V)= M/( V+ε) for the corresponding bias-corrected moments M M and V V. The weight update is then Wt←Wt−1−αη⋅G~tW_t← W_t-1-αη· G_t where η is the learning rate and α is a scale factor. Table 2: Optimizer-state memory requirement for an m×nm× n matrix. Optimizer State size Adam 2mn2mn SOAP 2mn+2m2+2n22mn+2m^2+2n^2 Muon mnmn ALSO mn+2mr+2nr+r2+2r+O(1)mn+2mr+2nr+~~r^2+2r+O(1) Clean (Ours) 2mr+2nr+2r2+4r+O(1)2mr+2nr+2r^2+4r+O(1) Clean avoids building any (mn)O(mn) memory while leveraging full-rank information and maintaining performance by (i) representing preconditioners with Nyström factors Ut,∗,Λt,∗U_t,*, _t,*, (i) operating on four complementary blocks, and (i) updating the stored low-rank moments when bases change. Memory therefore grows only linear with respect to model dimensions. Furthermore, our warm-start Nyström variant (see Alg. 3 in Appendix A) enables gradual basis evolution, improving stability during periodic subspace refreshes. As shown in Table 2, Clean is the only evaluated optimizer that scales linearly in state memory. Algorithm 1 Clean at iteration t. Colors mark the blocks of Eqn. (10): Main, A, B, C. Between refreshes, the previous factors are reused. 0: Gt∈ℝm×nG_t ^m× n, step size η, scale α, β1,β2,ϵ _1, _2,ε, rank r, refresh period T, and t≥1t≥ 1 1: if t−1modT=0t-1 T=0 then 2: Ut,L,Λt,L←Nys(β2Ut−1,LΛt−1,LUt−1,L⊤+(1−β2)GtGt⊤)U_t,L, _t,L ( _2U_t-1,L _t-1,LU_t-1,L +(1- _2)G_tG_t ) ⊳ Alg. 3 3: Ut,R,Λt,R←Nys(β2Ut−1,RΛt−1,RUt−1,R⊤+(1−β2)Gt⊤Gt)U_t,R, _t,R ( _2U_t-1,R _t-1,RU_t-1,R +(1- _2)G_t G_t ) ⊳ Alg. 3 4: Rotate M,MA,MB,vA,vBM,M_A,M_B,v_A,v_B into the new bases ⊳ Eqn. (14) & Eqn. (15) 5: Ut,L,Λt,L←Ut−1,L,Λt−1,LU_t,L, _t,L← U_t-1,L, _t-1,L 6: Ut,R,Λt,R←Ut−1,R,Λt−1,RU_t,R, _t,R← U_t-1,R, _t-1,R 7: St←Ut,LGtUt,R [rgb]0.1211,0.4648,0.707S_t← U_t,L TG_tU_t,R, Gt,A←PL⟂GtUt,R, [rgb]0.8398,0.4297,0.0781G_t,A← P_L G_tU_t,R, Gt,B←Ut,LGtPR⟂ [rgb]0.1719,0.5898,0.1719G_t,B← U_t,L TG_tP_R , Ct←PL⟂GtPR⟂ [rgb]0.4688,0.4688,0.4688C_t← P_L G_tP_R ⊳ Eqn. (10) 8: Mt [rgb]0.1211,0.4648,0.707M_t, Vt [rgb]0.1211,0.4648,0.707V_t, Mt,A [rgb]0.8398,0.4297,0.0781M_t,A, vt,A [rgb]0.8398,0.4297,0.0781v_t,A, Mt,B [rgb]0.1719,0.5898,0.1719M_t,B, vt,B [rgb]0.1719,0.5898,0.1719v_t,B, vt,C [rgb]0.4688,0.4688,0.4688v_t,C ⊳ Moments, Table 1 9: G~t←Ut,L(Mt,Vt)Ut,R+(Mt,A,vt,A)Ut,R+Ut,L(Mt,B,vt,B)+(Ct,vt,C) G_t← [rgb]0.1211,0.4648,0.707U_t,L\,N(M_t,V_t)\,U_t,R T+ [rgb]0.8398,0.4297,0.0781N(M_t,A,v_t,A)\,U_t,R T+ [rgb]0.1719,0.5898,0.1719U_t,L\,N(M_t,B,v_t,B)+ [rgb]0.4688,0.4688,0.4688N(C_t,v_t,C) ⊳ Eqn. (16) 10: Wt←Wt−1−αηG~tW_t← W_t-1-αη\, G_t 4 Q-Clean: Low-Precision Second-Order Optimization To further reduce memory overhead, we explore the stability of Clean under quantization, a technique that stores matrices as low-bit integers. In this section, we briefly introduce Q-Clean, which employs quantized optimizer states. More details are included in Appendix B. Following (Dettmers et al., 2021), block-wise quantization partitions a flattened matrix X into contiguous blocks of size B and quantizes each block independently. With an n-bit asymmetric (affine) quantizer, block i uses a scale sis_i and a zero-point ziz_i to represent a value in lower precision. Per-block scaling limits the effect of outliers to a single block rather than coarsening the quantization grid of the entire matrix (Dettmers et al., 2022). However, because optimizer states are continually updated, round-to-nearest scheme discards any update smaller than half a quantization step, so small but consistent updates never accumulate. To remove this bias, the linear quantizer can instead use stochastic rounding (Von Neumann and Goldstine, 1947; Gupta et al., 2015), which is unbiased in expectation, and therefore prevents small updates from vanishing over time. For the first-moment and projection matrices, we apply the block-wise linear asymmetric scheme mentioned above. In contrast, the second-moments are strictly non-negative and exhibit a highly skewed distribution with a large dynamic range (see Appendix B). Because linear quantization is suboptimal for such statistics, we instead adopt the non-linear quantization implementation from bitsandbytes (Dettmers et al., 2021), which is also used in the Adam-8bit optimizer. We quantize the first moments of all blocks listed in Table 1, the second moment of the Main, and the projection matrices in INT8. The second moments of blocks A and B are kept in full precision: each is a vector of size r, so its memory footprint is negligible and quantizing it would yield no meaningful savings. 5 Related Work Low-rank Projection Methods. A popular approach to reducing optimizer memory is to project gradients onto a low-rank subspace. GaLore (Zhao et al., 2024) introduced this strategy by performing gradient updates within a low-rank subspace defined by the singular value decomposition (SVD) of gradients. A family of works follows this line of optimizers (Chen et al., ; Zhu et al., 2025; Liang et al., 2024; He et al., 2024). While effective at saving memory, this one-sided projection can be unstable as the subspace drifts over time. LDAdam (Robert et al., 2024) and Subtrack++ (Rajabi et al., 2025) attempted to mitigate this by incorporating projection-aware updates and more robust subspace tracking mechanisms. However, these methods operate on one-sided approximations and do not capture the row-column geometry of the gradient covariances. Second-order Preconditioning. Another line of work focuses on capturing curvature information. K-FAC (Martens and Grosse, 2015) approximates the Fisher information matrix with Kronecker products, while Shampoo (Gupta et al., 2018) directly approximates the full-matrix AdaGrad (Duchi et al., 2011) preconditioner by maintaining separate row and column preconditioners. These balanced, two-sided preconditioners are powerful but require storing full-sized covariance matrices, making them memory- and compute-intensive. SOAP (Vyas et al., 2024) combines Shampoo’s preconditioning with first-order updates but still retains the costly full-matrix accumulators. Low-rank Second-order Methods. To achieve the benefits of second-order preconditioning without the prohibitive memory costs, recent methods have proposed efficient second-order optimizers. Alice (Gong et al., 2026) provides a low-rank extension that converts general full-rank optimizers into memory-efficient approximations by relying on structured Fisher information matrix approximations. Building on this, ALSO (Romashov et al., 2026), extends toward a two-sided approximation of SOAP (Vyas et al., 2024), while compensating for the discarded residual blocks outside the core low-rank projection. Another direction leverages the varying importance of eigensubspaces; for instance, COSMOS (Liu et al., 2025) introduces a hybrid approach that applies memory-intensive SOAP updates exclusively to the leading eigensubspace of the gradient matrix, while utilizing a computationally cheaper, Muon-like update (Jordan et al., 2024) for the residual subspace. Low Precision Training with Quantization. Quantization improves memory efficiency by representing data in reduced-precision formats. These techniques broadly fall into Post-Training Quantization (PTQ), where pre-trained models are quantized without additional training (Ashkboos et al., 2025; Dettmers et al., 2023; Ashkboos et al., 2024), and Quantization-Aware Training (QAT), which integrates quantization directly into the training loop. A major challenge in QAT is training instability. Inspired by Izmailov et al. (2018), SWALP (Yang et al., 2019) mitigates this using stochastic weight averaging, though it incurs significant memory overhead when applied to large foundation models. Q-Galore (Zhang et al., 2024) substantially reduces memory consumption by integrating full-model quantization with the low-rank gradient projection of GaLore (Zhao et al., 2024). Similarly, 4-bit Shampoo (Wang et al., 2024) quantizes the left and right eigenvector matrices and applies Björck orthonormalization to restore orthogonality after quantization. 6 Experiments Figure 4: Validation perplexity recorded during the training of a LLaMA-1.3B architecture across various optimization methods. Setup. We evaluate Clean and Q-Clean on pre-training LLaMA-based architectures (Touvron et al., 2023) with 350M and 1.3B parameters on the C4 dataset (Raffel et al., 2020). All models are trained from scratch with a sequence length of 1024 and an effective batch size of 4096 (≈4.2≈ 4.2M tokens per step), identical across optimizers. Following the Chinchilla rule (Hoffmann et al., 2022) of ≈20≈ 20 training tokens per parameter, we train for 1750 steps (7.37.3B tokens) at 350M and 5000 steps (21.021.0B tokens) at 1.3B. We use a cosine schedule with 10%10\% linear warmup decaying to 0.1×0.1× the peak learning rate, weight decay 0.10.1, gradient clipping at 1.01.0, β1=β2=0.95 _1= _2=0.95, and bfloat16 parameters and gradients. Training is carried out on nodes of up to four NVIDIA H100 80GB GPUs. Baselines and tuning. We compare against AdamW (Kingma and Ba, 2014), SOAP (Vyas et al., 2024), Muon (Jordan et al., 2024), and ALSO (Romashov et al., 2026). At 350M we sweep three learning rates per method and report the best run in Table 3. Due to the prohibitive cost of a full sweep at 1.3B, we performed a shortened sweep and resumed only the best configurations to 5,000 steps. Clean, Q-Clean and ALSO all use rank r=256r=256 at 350M and r=512r=512 at 1.3B, and all preconditioned methods refresh their preconditioner every T=10T=10 steps. Table 3: Final validation perplexity (↓ ) and optimizer memory when pre-training LLaMA models of different sizes. Clean achieves competitive perplexity compared to full SOAP, while using less memory than AdamW. Clean, Q-Clean and ALSO use rank r=256r=256 at 350M and r=512r=512 at 1.3B. Excluding SOAP, the best results are shown in bold and the second best are underlined. Methods LLaMA-350M LLaMA-1.3B Perplexity Opt. Mem. Peak Mem. Perplexity Opt. Mem. Peak Mem. AdamW (Kingma and Ba, 2014) 19.63 1.40G 47.93G 12.12 5.11G 54.13G SOAP (Vyas et al., 2024) 15.12 7.64G 54.25G 11.52 13.62G 79.02G Muon (Jordan et al., 2024) 16.97 0.83G 47.36G 11.78 2.80G 51.83G ALSO (Romashov et al., 2026) 15.36 1.30G 47.88G 12.19 4.71G 53.82G Clean (Ours) 15.87 1.25G 47.68G 11.65 4.48G 53.44G Q-Clean (Ours) 16.44 0.38G 46.82G 12.10 1.28G 50.24G Results. Table 3 summarizes the final results, demonstrating that Clean achieves competitive performance compared to SOAP, at a fraction of its cost. Q-Clean cuts the optimizer state memory consumption by more than 50% compared to Muon, and still delivers results on par with the full-rank methods. Fig. 4 plots validation perplexity throughout training for LLaMA-1.3B, where Clean follows SOAP closely and presents on-par results. Efficiency. Clean’s memory advantage grows with scale (Figs. 7 and 8) because its memory complexity scales linearly with matrix dimensions, at a fixed rank r, whereas every baseline carries at least one quadratic term (Table 2). Furthermore, Fig. 7 demonstrates a reduction in per-parameter memory cost as model size increases. See Appendix D for more details. (a) Validation loss against wall-clock time. (b) Steps to reach AdamW’s final loss. Figure 5: Convergence efficiency on pre-training LLaMA-1.3B, all runs on four H100 GPUs with the same batch size, schedule and token budget. (a) Wall-clock: Clean reaches AdamW’s final validation loss in 6464 of the 8181 hours of AdamW requirements, demonstrating 26% speedup. (b) Token efficiency: Clean reaches that loss in 77%77\% of the training budget and Q-Clean in 87%87\%. To evaluate sample efficiency and convergence, Fig. 5 compares the methods across time and token budgets rather than focusing solely on the final loss. Clean reaches AdamW’s final performance in 64 hours instead of 81, while simultaneously reducing the optimizer memory footprint. It also only consumes 77% of the overall token budget, with Q-Clean requiring 87%. Ablations. To better investigate the sensitivity of Clean, we performed ablation studies on the LLaMA-350M pre-training, focusing on two key settings: the projection rank r and the preconditioner update interval T. In each experiment, we vary one parameter while keeping the other fixed at our default configuration (T=10T=10 and r=256r=256). First, Fig. 9(a) shows that the validation perplexity improves as the rank increases up to 256. While heavily compressed ranks (r=32,64r=32,64) fail to catch up to the baselines, setting r=128r=128 achieves competitive results, and increasing the rank further to r=512r=512 yields no additional gains. Second, Fig. 9(b) demonstrates that T=10T=10 provides an optimal balance; while refreshing the preconditioners either too frequently (T=1T=1) or at large intervals (T=100T=100) negatively impacts the final performance. 7 Conclusion We introduce Clean, a memory-efficient second-order optimizer that reduces the memory complexity of two-sided preconditioning from quadratic to linear. By leveraging randomized Nyström sketching, Clean accurately approximates left and right preconditioners, while uniquely reintegrating off-subspace residual blocks to preserve essential training signals. We also develop Q-Clean, a low-precision variant that aggressively compresses optimizer states. Empirically, both Clean and Q-Clean achieve validation perplexities that are highly competitive with full-rank methods like SOAP, yet they operate with a smaller memory footprint than standard AdamW. Crucially, this aggressive memory reduction uniquely enables the pre-training of a 13B parameter architecture on a single 80GB GPU, and increases the convergence speed by 26% compared to AdamW. We are optimistic that these linear-scaling curvature approximations will encourage broader exploration into accessible, memory-efficient optimization frameworks. Acknowledgment Sirisha Rambhatla and Sahar Rajabi acknowledge support from the Natural Sciences and Engineering Research Council of Canada (NSERC) Discovery Grant RGPIN-2022-03512 and the Val O’Donovan Chair endowment. Guillaume Rabusseau and Beheshteh T. Rakhshan acknowledge support from the CIFAR AI Chair Program. This work was made possible in part through computing resources provided by the Digital Research Alliance of Canada and Mila (mila.quebec). AI Use Statement In this work, we used generative AI tools to provide feedback on our research methodology and to assist in implementing methods. We did not use generative AI for other tasks requiring disclosure, such as generating synthetic datasets, help develop theoretical models or conceptual frameworks, propose or refine hypothesis, design research methodology, assist in the writing of proofs, provide critical ingredients for proving mathematical claims, assist with translation, clean and reformat dataset, support qualitative and thematic data analysis, developing theoretical models, formulating or proving mathematical claims, or interpreting the final results. Additionally, we utilized generative AI for several tasks with recommended disclosure: analyzing existing literature, creating scientific figures and images, and revising the paper draft to improve wording, flow, and overall readability. We have thoroughly reviewed and verified all AI-assisted work. We take full responsibility for the final content of this work, including all text, claims, code, and visual artifacts produced with the aid of generative AI. References Anil et al. (2020) R. Anil, V. Gupta, T. Koren, K. Regan, and Y. Singer Scalable second order optimization for deep learning. arXiv preprint arXiv:2002.09018. Cited by: §2.1. Ashkboos et al. (2024) S. Ashkboos, A. Mohtashami, M. L. Croci, B. Li, P. Cameron, M. Jaggi, D. Alistarh, T. Hoefler, and J. Hensman Quarot: outlier-free 4-bit inference in rotated llms. Advances in Neural Information Processing Systems 37, p. 100213–100240. Cited by: §5. Ashkboos et al. (2025) S. Ashkboos, M. Nikdan, S. Tabesh, R. L. Castro, T. Hoefler, and D. Alistarh Halo: hadamard-assisted lower-precision optimization for llms. arXiv preprint arXiv:2501.02625. Cited by: §5. [4] X. Chen, K. Feng, C. Li, X. Lai, X. Yue, Y. Yuan, and G. Wang Fira: can we achieve full-rank training of llms under low-rank constraint?, 2024. URL https://arxiv. org/abs/2410.01623. Cited by: §5. Dettmers et al. (2022) T. Dettmers, M. Lewis, Y. Belkada, and L. Zettlemoyer Gpt3. int8 (): 8-bit matrix multiplication for transformers at scale. Advances in neural information processing systems 35, p. 30318–30332. Cited by: §B.1, §4. Dettmers et al. (2021) T. Dettmers, M. Lewis, S. Shleifer, and L. Zettlemoyer 8-bit optimizers via block-wise quantization. arXiv preprint arXiv:2110.02861. Cited by: §B.1, §B.1, §4, §4. Dettmers et al. (2023) T. Dettmers, A. Pagnoni, A. Holtzman, and L. Zettlemoyer Qlora: efficient finetuning of quantized llms. Advances in neural information processing systems 36, p. 10088–10115. Cited by: §B.1, §5. Duchi et al. (2011) J. Duchi, E. Hazan, and Y. Singer Adaptive subgradient methods for online learning and stochastic optimization.. Journal of machine learning research 12 (7). Cited by: §2.1, §2.1, §5. Frangella et al. (2023) Z. Frangella, J. A. Tropp, and M. Udell Randomized nyström preconditioning. SIAM Journal on Matrix Analysis and Applications 44 (2), p. 718–752. Cited by: Appendix A, §2.2, §3. Gittens and Mahoney (2016) A. Gittens and M. W. Mahoney Revisiting the nyström method for improved large-scale machine learning. The Journal of Machine Learning Research 17 (1), p. 3977–4041. Cited by: Appendix A, §2.2, §2.2. Gong et al. (2026) W. Gong, M. Scetbon, C. Ma, and E. Meeds Towards efficient optimizer design for LLM via structured fisher approximation with a low-rank extension. In The Fourteenth International Conference on Learning Representations, External Links: Link Cited by: §1, §5. Gupta et al. (2015) S. Gupta, A. Agrawal, K. Gopalakrishnan, and P. Narayanan Deep learning with limited numerical precision. In International conference on machine learning, p. 1737–1746. Cited by: §B.2, §4. Gupta et al. (2018) V. Gupta, T. Koren, and Y. Singer Shampoo: preconditioned stochastic tensor optimization. In International Conference on Machine Learning, p. 1842–1850. Cited by: §1, §1, §2.1, §3, §5. Halko et al. (2011) N. Halko, P. Martinsson, and J. A. Tropp Finding structure with randomness: probabilistic algorithms for constructing approximate matrix decompositions. SIAM review 53 (2), p. 217–288. Cited by: §3. He et al. (2024) Y. He, P. Li, Y. Hu, C. Chen, and K. Yuan Subspace optimization for large language models with convergence guarantees. arXiv preprint arXiv:2410.11289. Cited by: §5. Hoffmann et al. (2022) J. Hoffmann, S. Borgeaud, A. Mensch, E. Buchatskaya, T. Cai, E. Rutherford, D. d. L. Casas, L. A. Hendricks, J. Welbl, A. Clark, et al. Training compute-optimal large language models. arXiv preprint arXiv:2203.15556. Cited by: §6. Izmailov et al. (2018) P. Izmailov, D. Podoprikhin, T. Garipov, D. Vetrov, and A. G. Wilson Averaging weights leads to wider optima and better generalization. arXiv preprint arXiv:1803.05407. Cited by: §5. Jordan et al. (2024) K. Jordan, Y. Jin, V. Boza, J. You, F. Cesista, L. Newhouse, and J. Bernstein Muon: an optimizer for hidden layers in neural networks. External Links: Link Cited by: §1, §5, §6, Table 3. Kingma and Ba (2014) D. P. Kingma and J. Ba Adam: a method for stochastic optimization. arXiv preprint arXiv:1412.6980. Cited by: §1, §2.1, §6, Table 3. Liang et al. (2024) K. Liang, B. Liu, L. Chen, and Q. Liu Memory-efficient llm training with online subspace descent. Advances in Neural Information Processing Systems 37, p. 64412–64432. Cited by: §5. Liu et al. (2025) L. Liu, Z. Xu, Z. Zhang, H. Kang, Z. Li, C. Liang, W. Chen, and T. Zhao Cosmos: a hybrid adaptive optimizer for memory-efficient training of llms. arXiv preprint arXiv:2502.17410. Cited by: §1, §5. Martens and Grosse (2015) J. Martens and R. Grosse Optimizing neural networks with kronecker-factored approximate curvature. In International conference on machine learning, p. 2408–2417. Cited by: §1, §5. Morwani et al. (2024) D. Morwani, I. Shapira, N. Vyas, E. Malach, S. Kakade, and L. Janson A new perspective on shampoo’s preconditioner. arXiv preprint arXiv:2406.17748. Cited by: §3. Raffel et al. (2020) C. Raffel, N. Shazeer, A. Roberts, K. Lee, S. Narang, M. Matena, Y. Zhou, W. Li, and P. J. Liu Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research 21 (140), p. 1–67. Cited by: §6. Rajabi et al. (2025) S. Rajabi, N. Nonta, and S. Rambhatla SubTrack++: gradient subspace tracking for scalable llm training. arXiv preprint arXiv:2502.01586. Cited by: §5. Robert et al. (2024) T. Robert, M. Safaryan, I. Modoranu, and D. Alistarh Ldadam: adaptive optimization from low-dimensional gradient statistics. arXiv preprint arXiv:2410.16103. Cited by: §5. Romashov et al. (2026) F. Romashov, N. Yudin, and M. Rakhuba Two-sided low-rank SOAP for efficient LLM training. In ICML’26 workshop on CoLoRAI - The 2nd Workshop on Connecting Low-rank Representations in AI, External Links: Link Cited by: Appendix D, Appendix D, §1, §3, §5, §6, Table 3. Shazeer and Stern (2018) N. Shazeer and M. Stern Adafactor: adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning, p. 4596–4604. Cited by: §1. Touvron et al. (2023) H. Touvron, L. Martin, K. Stone, P. Albert, A. Almahairi, Y. Babaei, N. Bashlykov, S. Batra, P. Bhargava, S. Bhosale, et al. Llama 2: open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288. Cited by: §6. Tropp et al. (2017) J. A. Tropp, A. Yurtsever, M. Udell, and V. Cevher Fixed-rank approximation of a positive-semidefinite matrix from streaming data. Advances in Neural Information Processing Systems 30. Cited by: Appendix A, §2.2. Von Neumann and Goldstine (1947) J. Von Neumann and H. H. Goldstine Numerical inverting of matrices of high order. Cited by: §B.2, §4. Vyas et al. (2024) N. Vyas, D. Morwani, R. Zhao, M. Kwun, I. Shapira, D. Brandfonbrener, L. Janson, and S. Kakade Soap: improving and stabilizing shampoo using adam. arXiv preprint arXiv:2409.11321. Cited by: Appendix C, §1, §1, §2.1, §3, §3, §5, §5, §6, Table 3. Wang et al. (2024) S. Wang, P. Zhou, J. Li, and H. Huang 4-bit shampoo for memory-efficient network training. Advances in Neural Information Processing Systems 37, p. 126997–127029. Cited by: §5. Yang et al. (2019) G. Yang, T. Zhang, P. Kirichenko, J. Bai, A. G. Wilson, and C. De Sa Swalp: stochastic weight averaging in low precision training. In International Conference on Machine Learning, p. 7015–7024. Cited by: §5. Zhang et al. (2024) Z. Zhang, A. Jaiswal, L. Yin, S. Liu, J. Zhao, Y. Tian, and Z. Wang Q-galore: quantized galore with int4 projection and layer-adaptive low-rank gradients. arXiv preprint arXiv:2407.08296. Cited by: §B.2, §5. Zhao et al. (2024) J. Zhao, Z. Zhang, B. Chen, Z. Wang, A. Anandkumar, and Y. Tian Galore: memory-efficient llm training by gradient low-rank projection. arXiv preprint arXiv:2403.03507. Cited by: §3, §5, §5. Zhu et al. (2025) H. Zhu, Z. Zhang, W. Cong, X. Liu, S. Park, V. Chandra, B. Long, D. Z. Pan, Z. Wang, and J. Lee APOLLO: SGD-like memory, adamw-level performance. In Eighth Conference on Machine Learning and Systems, External Links: Link Cited by: §5. Appendix A Nyström Algorithm Algorithm 2 Randomized Nyström Approximation (RandNystromApprox) 1: Input: Matrix A∈ℝm×mA ^m× m (PSD), rank size r≤mr≤ m, Ω∈ℝm×r ^m× r. 2: Ω∼(0,I)∈ℝm×r (0,I) ^m× r random Gaussian matrix 3: Ω←QR(Ω,0) ( ,0) thin QR decomposition 4: Y←AΩY← A 5: ν←eps(norm(Y,2))ν (norm(Y,2)) compute shift 6: Yν←Y+νΩY_ν← Y+ν add shift for stability 7: B←ΩYνB← TY_ν 8: C←CHOL((B+B)/2)C ((B+B T)/2) Force symmetry and Cholesky decomposition 9: B←Yν/CB← Y_ν/C triangular solve 10: [U,Σ,∼]=SVD(B,0)[U, , ]=SVD(B,0) thin SVD 11: Λ^←max0,Σ2−νI \0, ^2-ν I\ remove shift, compute eigs 12: Output: [U,Λ^][U, ] Nyström approximations are well-suited for optimizer states due to their computational efficiency and solid theoretical foundations (Gittens and Mahoney, 2016). While a direct implementation using Eqn. (4) suffers from numerical instability, Alg. 2 introduces a robust, numerically stable alternative for generating randomized Nyström approximations, as proposed in Tropp et al. (2017) and Frangella et al. (2023). In this paper, rather than sampling a purely random matrix at each step, we warm-start the sketch using the eigenvectors from the previous step (UprevU_prev), augmented with p columns of fresh Gaussian noise to continuously explore new gradient directions. To further improve approximation accuracy, we perform q steps of power iteration with periodic orthogonalization. Finally, we apply sign alignment to the resulting eigenvectors to prevent abrupt subspace flips, ensuring temporal consistency across consecutive optimization steps. The pseudo-code of this updated method is presented in Alg. 3. Algorithm 3 Warm-started Randomized Nyström with Subspace Iteration 0: A∈ℝm×mA ^m× m (PSD), target rank r, oversampling p, power-iterations passes q, warm start Uprev∈ℝm×rU_prev ^m× r (or ∅ ) 1: if Uprev≠∅U_prev≠ then 2: Ψ∼(0,I)∈ℝm×p (0,I) ^m× p; Ω←[Uprev,Ψ] ←[\,U_prev,\, \,] warm start + fresh exploration block 3: else 4: Ω∼(0,I)∈ℝm×(r+p) (0,I) ^m×(r+p) 5: Ω←QR(Ω,0) ( ,0) 6: for j=1j=1 to q do 7: Y←AΩY← A 8: Ω←QR(Y,0) (Y,0) re-orthonormalize every pass 9: Y←AΩY← A final sketch must match the final Ω 10: ν←eps(norm(Y,2))ν (norm(Y,2)) 11: Yν←Y+νΩY_ν← Y+ν 12: B←ΩYνB← TY_ν 13: C←CHOL((B+B)/2)C ((B+B T)/2) 14: [U,Σ,∼]←SVD(Yν/C,0)[U, , ] (Y_ν/C,0) 15: Λ^←max0,Σ2−νI ← \0, ^2-ν I\ 16: U←U:,1:rU← U_:,1:r, Λ^←Λ^1:r,1:r ← _1:r,1:r discard oversampling 17: si←sign(uiuprev,i)s_i (u_i Tu_prev,i); U:,i←siU:,iU_:,i← s_iU_:,i sign alignment 18: Output: [U,Λ^][U, ] Appendix B Quantization Details In this section, we explain the Q-Clean optimizer in more detail. B.1 Block-wise n-bit Quantization Block-wise quantization was originally introduced by (Dettmers et al., 2021) for compressing optimizer states during training, and was later adapted by (Dettmers et al., 2023) to frozen base-model weights for memory-efficient LLM fine-tuning. We adopt the same block-wise principle but use a uniform affine quantizer, with an asymmetric formulation (scale and zero-point) for most matrices and a symmetric formulation for the second moment, as detailed below. The block-wise scheme partitions a matrix into contiguous blocks and computes an independent scale s (and, in the asymmetric case, a zero-point z) for each block. The motivation is robustness to outliers: with a single per matrix scale, even one anomalously large value forces every quantization step to be coarse, mapping most typical values to only a handful of integer codes and losing most of the resolution (Dettmers et al., 2022). A per-block scale isolates this effect: only the block containing the outlier loses precision, while every other block keeps a fine-grained scale fitted to its own local data. Asymmetric Quantization (Default). Given a matrix X∈ℝd1×d2X ^d_1× d_2, we flatten and partition it into N=⌈d1d2/B⌉N= d_1d_2/B contiguous blocks X(i)i=1N\X^(i)\_i=1^N of size at most B (the last block may be smaller). Each block is quantized independently with a per-block scale sis_i and zero-point ziz_i computed from the block’s full range: si=max(X(i))−min(X(i))2n−1,zi=−⌊min(X(i))/si⌉s_i= (X^(i))- (X^(i))2^n-1, z_i=- (X^(i))/s_i (17) Then, the quantization and dequantization maps are Xq(i)=Quantnasym(X(i),si,zi)=clamp(⌊X(i)si⌉+zi, 0, 2n−1),X^(i)=Dequantnasym(Xq(i),si,zi)=si(Xq(i)−zi)≈X(i) splitX_q^(i)&=Quant_n^asym(X^(i),s_i,z_i)=clamp\! ( X^(i)s_i +z_i,\;0,\;2^n-1 ),\\ X^(i)&=Dequant_n^asym(X_q^(i),s_i,z_i)=s_i (X_q^(i)-z_i )\;≈\;X^(i) split (18) where ⌊⋅⌉ · denotes round-to-nearest and n is the number of bits. We apply this affine form with the stochastic-rounding explained in section B.2 to the projection matrices (Ut,LU_t,L and Ut,RU_t,R) and to all first-moment matrices (MtM_t, Mt,AM_t,A, and Mt,BM_t,B). However, we do not use this quantization for the second-moment matrices, because their strictly non-negative, highly skewed distributions and large dynamic ranges make linear methods suboptimal. Instead, we apply a non-linear symmetric quantization approach for the main second moment. The second moments of blocks A and B remain in full precision; as simple size-r vectors, their memory footprint is already negligible and quantizing them offers no meaningful savings. Non-linear Quantization for the Second-moment. Because the second moment statistics are strictly non-negative (X(i)≥0X^(i)≥ 0) and exhibit a highly skewed distribution with a large dynamic range, standard linear quantization allocates insufficient resolution to the majority of values concentrated near zero. To address this, we utilize the non-linear quantization scheme from bitsandbytes (Dettmers et al., 2021). Rather than spacing bins equally, this method employs a fixed non-linear codebook C consisting of 256 quantization levels (for 8-bit representation). These levels are dynamically spaced to provide high precision for small-magnitude values while remaining capable of capturing large outliers. For a given block, we first compute the maximum value to serve as the block-wise scaling factor: si=max(X(i))s_i= (X^(i)) (19) The elements of the block are normalized by sis_i to lie in the range [0,1][0,1]. Each normalized value is then mapped to the nearest entry in the codebook C. The quantization and dequantization mappings are formalized as: Xq(i)=Quant8-bitNL(X(i),si)=argminc∈|X(i)si−c|,X^(i)=Dequant8-bitNL(Xq(i),si)=si⋅Xq(i), splitX_q^(i)&=Quant_8-bit^NL(X^(i),s_i)= _c | X^(i)s_i-c |,\\ X^(i)&=Dequant_8-bit^NL(X_q^(i),s_i)=s_i· X_q^(i), split (20) where Xq(i)X_q^(i) is stored in memory as an 8-bit integer corresponding to the index of the matched codebook entry c. This formulation eliminates the need for a zero-point, requiring only the per-block scale sis_i as metadata, while perfectly matching the non-linear, non-negative distribution of the optimizer’s second moments. B.2 Maintaining High-precision Training with Stochastic Rounding Accumulated Quantization To reduce memory footprint, we store optimizer states as matrices in low precision (INT8) throughout the forward and backward passes. Doing so, however, introduces a subtle but fatal failure mode: under round-to-nearest (RTN), if the update magnitude is not large enough, the effect will be mitigated by RTN. Following Q-GaLore (Zhang et al., 2024), we apply stochastic rounding (SR) (Von Neumann and Goldstine, 1947; Gupta et al., 2015) in our asymmetric quantization. A floating-point value X lying between two adjacent quantization levels ⌊X⌋ X and ⌈X⌉ X is rounded up to ⌈X⌉ X with probability X−⌊X⌋X- X and rounded down to ⌊X⌋ X otherwise. This rule is unbiased, i.e. (Xq)=XE(X_q)=X. More concretely, the two rounding schemes, namely rounding-to-nearest (RTN) and stochastic rounding, are defined as follows, respectively: ⌊X⌉ X =⌊X⌋if ⌊X⌋≤X<⌊X⌋+ϵ2,⌊X⌋+ϵif ⌊X⌋+ϵ2≤X≤⌊X⌋+ϵ, = cases X &if X ≤ X< X + ε2,\\[2.0pt] X +ε&if X + ε2≤ X≤ X +ε, cases (21) round(X)=⌊X⌋with probability ⌈X⌉−X,⌈X⌉with probability X−⌊X⌋, (X)= cases X &with probability X -X,\\[2.0pt] X &with probability X- X , cases (22) where ϵ=2−FLε=2^-FL and FLFL denotes the length of the fractional part. B.3 Quantizer Choice We tracked every matrix family W (weights), G (gradients), M and V (Adam moments), L and R (SOAP’s accumulated covariance matrices), and QLQ_L and QRQ_R (eigenbasis of L and R) at every layer of the LLaMA-350M model while optimizing via SOAP (see Alg. 4). For each family, we then pool all captured matrices and report empirical densities of three quantities: scale (log10max|x| _10 |x|), dynamic range (log10(max|x|/median|x|)) ( _10 ( |x|/median|x| ) ), and the non-excess kurtosis (log10κ _10κ), with κ=[(x−μ)4]/σ4κ=E[(x-μ)^4]/σ^4. As shown in Figure 6, we evaluate the statistics using two perspectives based on the quantization strategy. The per-matrix view computes a single scalar for an entire matrix, highlighting differences between matrices; this applies when sharing one quantization scale across a matrix. Conversely, the per-block view divides matrices into blocks, computing and pooling statistics for each block individually. Together, these views motivate our quantization strategy. In the per-matrix view, overall matrices (V, L, R, etc.) exhibit extreme distributional probabilities; a global linear scale stretches to cover these heavy tails, wasting resolution and causing high error. Moving to a block-wise approach enhances the tiles by isolating these extreme values, which lightens the tails across all families. For most matrices (W, G, M, L, R, QLQ_L, QRQ_R), this shift makes them suitable for asymmetric and linear quantization. However, even after block-wise quantization, the second-moment matrix (V) still shows heavy tails, and linear quantization leads to divergence on V, whereas it succeeds for the other matrices. Therefore, we use a symmetric and non-linear quantizer for V and retain linear block-wise codes for the rest in our low-rank proposed method. (a) log10κ _10^κ per block (κ=[(x−μ)4]/σ4κ=E[(x-μ)^4]/σ^4). (b) log10max(|x|)/median(|x|) _10^max(|x|)/median(|x|) per block. (c) log10max(|x|) _10 (|x|) per block. (d) log10κ _10^κ per matrix (κ=[(x−μ)4]/σ4κ=E[(x-μ)^4]/σ^4). (e) log10max(|x|)/median(|x|) _10^max(|x|)/median(|x|) per matrix. (f) log10max(|x|) _10 (|x|) per matrix . Figure 6: The plots compare non-excess kurtosis, dynamic range, and scale under two quantization granularities: per-block (a–c) and per-matrix (d–f). The shift from a per-matrix to a per-block approach significantly reduces the heavy-tailed distributions for most matrices, rendering them suitable for linear quantization. However, the second moment matrix (V, in purple) remains a prominent outlier, exhibiting persistent heavy tails even in the per-block view, which motivates the use of a non-linear quantizer specifically for V. Appendix C SOAP Algorithm 4 presents the pseudocode for the SOAP optimizer (Vyas et al., 2024). As shown, the full L and R covariance matrices are constructed and stored, while the Adam optimizer is applied to the gradient projected onto their eigenbasis. Although SOAP demonstrates accelerated convergence compared to other standard optimizers, it introduces prohibitive memory requirements by simultaneously maintaining both the covariance matrices and the full Adam moments. 1: Sample batch BtB_t. 2: G∈ℝm×n←−∇WϕBt(Wt)G ^m× n←- _W _B_t(W_t) 3: G′←QLTGQRG ← Q_L^TGQ_R 4: M←β1M+(1−β1)GM← _1M+(1- _1)G 5: M′←QLTMQRM ← Q_L^TMQ_R 6: Now we “run” Adam on G′G 7: V←β2V+(1−β2)(G′⊙G′)V← _2V+(1- _2)(G G ) Elementwise multiplication 8: N′←M′V^t+ϵN ← M V_t+ε Elementwise division and square root 9: Now that we have preconditioned by Adam in the rotated space, we go back to the original space. 10: N←QLN′QRTN← Q_LN Q_R^T 11: W←W−ηNW← W-η N 12: End of gradient step, we now update L and R and possibly also QLQ_L and QRQ_R. 13: L←β2L+(1−β2)GGTL← _2L+(1- _2)G^T 14: R←β2R+(1−β2)GTGR← _2R+(1- _2)G^TG 15: if t % f == 0 then 16: QL←Eigenvectors(L,QL)Q_L← Eigenvectors(L,Q_L) 17: QR←Eigenvectors(R,QR)Q_R← Eigenvectors(R,Q_R) Algorithm 4 SOAP Appendix D Linear Memory Discussion The linear memory complexity of Clean and Q-Clean with respect to model dimensions makes them exceptionally scalable. As demonstrated in Figure 7 (left), Clean and Q-Clean are the only optimizers that allow pre-training a LLaMA-13B architecture on a single 80 GB H100 GPU. Notably, the other low-rank second-order baseline, ALSO (Romashov et al., 2026), requires more than 80 GB of memory even at rank r=32r=32. This is due to the fact that Clean keeps no matrix with dimensions of order (mn)O(mn) or (m2+n2)O(m^2+n^2), while all other optimizers have at least one (mn)O(mn) component. Furthermore, Figure 7 (right) illustrates the cost per parameter for different optimizers. Remarkably, our methods are the only ones that reduce this per-parameter cost compared to the baselines. To ensure a realistic setup, we scale the rank proportionally with the model dimension, setting r=d/8r=d/8, where d=min(m,n)d= (m,n) is the smaller dimension of each weight matrix. This trend is highly promising in terms of democratizing training of much larger foundation models on more accessible hardware. Figure 7: Left: total memory for pre-training LLaMA-13B on a single 80 GB H100 as a function of rank; the shaded region is infeasible for all of our baselines at any rank. Right: optimizer state per parameter across model sizes at r≈d/8r≈ d/8. Every baseline carries an mnmn term (Table 2) and so stays flat or rises; Clean and Q-Clean scale as r(m+n)r(m+n) and are the only methods whose per-parameter cost decreases with model size. Additionally, Figure 8 illustrates the effect of varying ranks across different model sizes for Clean, Q-Clean, and the low-rank second-order baseline, ALSO (Romashov et al., 2026). Figure 8 (left) demonstrates that for a fixed rank of r=32r=32, the memory savings gap between our methods and the baseline increases as the model size scales. This improvement stems from the omission of quadratic memory terms, allowing our methods to operate at a significantly lower cost. Furthermore, Figure 8 (right) compares how the memory consumption of ALSO scales relative to our methods across different ranks and model dimensions. As expected, the memory efficiency gap widens as the model size increases. Figure 8: Measured optimizer state on LLaMA. Baselines carry an mnmn term; ours replaces it with r(r+2)r(r+2), so our state is smaller whenever mn>r(r+2)mn>r(r+2). The margin widens with scale because mnmn grows with the model (2.82.8M to 70.870.8M from 350M to 13B) while r(r+2)r(r+2) is fixed by the rank. Left: the shaded gap is the omitted mnmn term at r=32r=32. Right: reduction over ALSO, 15.7×15.7× at r=32r=32 and 3.8×3.8× at r=512r=512 on 13B. Batch 1, seq 1024, bf16 state; filled markers measured, hollow closed-form where the run exceeded 80 GB in ALSO. Appendix E Ablating Hyperparameters Here we provide the supporting figures for the hyperparameter ablation studies discussed in section 6. Fig. 9(a) presents the learning curves and validation perplexities for varying projection ranks (r) on the LLaMA-350M pre-training task. Figure 9(b) details the performance across different preconditioner update intervals (T). These figures illustrate the full training trajectories, confirming the optimal configurations of r=256r=256 and T=10T=10 established in our primary analysis. (a) Rank r in projection matrices. (b) Preconditioner update frequency T. Figure 9: Ablations of Clean on pre-training LLaMA-350M. (a) Rank r: Performance improves as rank increases up to 128 with marginal improvements afterward. (b) Update interval T: T=10T=10 is optimal; updating too often (T=1T=1) or too rarely (T=100T=100) negatively impacts performance. Appendix F Hyperparameters Tables 4 and 6 list the hyperparameters of all runs. SOAP, Clean and Q-Clean optimize all parameters themselves: 1-D parameters receive Adam updates with the method’s own (β1,β2)( _1, _2). Q-Clean quantizes the eigenbases and all first moments, including the M, MAM_ A, and MBM_ B, with asymmetric per-group linear quantization and stochastic rounding, and the second moments with block-wise dynamic quantization; the eigenvalues and the pooled second moments stay in original data type. Table 4: Model architectures and training setup shared by all optimizers. Gradient accumulation makes up the global batch. Peak memory in Table 3 is measured on a single GPU at this micro-batch size, without activation checkpointing. LLaMA-350M LLaMA-1.3B Architecture Parameters 368M 1.34B Hidden dimension 1024 2048 Intermediate dimension 2736 5461 Attention heads 16 32 Layers 24 24 Vocabulary size 32000 (LLaMA tokenizer) Training Dataset C4 (English) Sequence length 1024 Batch size (sequences) 4096 Micro-batch size 32 16 Training steps 1750 5000 Warmup steps 175 500 Learning-rate schedule linear warmup, cosine decay to 0.1×0.1× peak Weight decay (decoupled, all parameters) 0.1 Gradient clipping (global norm) 1.0 Adam ϵε 10−810^-8 Precision bf16 (weights, gradients, activations) Evaluation C4 validation, ≈ 10M tokens Table 5: Learning rates swept at 350M. Three values per optimizer; the selected value, used in Table 6 and reported in Table 3, is shown in bold. Method Learning rate AdamW 7×10−47× 10^-4 1.4×−1.4× 10^-3 2.8×10−32.8× 10^-3 SOAP 7×10−47× 10^-4 1.4×10−31.4× 10^-3 2.8×−2.8× 10^-3 Muon 2×10−32× 10^-3 6×10−36× 10^-3 1.0×−1.0× 10^-2 ALSO 1.4×10−31.4× 10^-3 4.4×−4.4× 10^-3 1.4×10−21.4× 10^-2 Clean (Ours) 1.4×10−31.4× 10^-3 4.4×−4.4× 10^-3 1.4×10−21.4× 10^-2 Q-Clean (Ours) 1.4×10−31.4× 10^-3 4.4×−4.4× 10^-3 1.4×10−21.4× 10^-2 Table 6: Optimizer-specific hyperparameters; learning rates are peak values of the schedule in Table 4. Method Hyperparameter LLaMA-350M LLaMA-1.3B AdamW Learning rate 1.4×10−31.4× 10^-3 1.4×10−31.4× 10^-3 (β1,β2)( _1, _2) (0.95,0.95)(0.95,0.95) SOAP Learning rate 2.8×10−32.8× 10^-3 2.8×10−32.8× 10^-3 (β1,β2)( _1, _2) (0.95,0.95)(0.95,0.95) Shampoo β 0.95 Preconditioning frequency 10 Muon Learning rate 1.0×10−21.0× 10^-2 1.4×10−21.4× 10^-2 Momentum (Nesterov) 0.9 Newton–Schulz iterations 5 ALSO Learning rate 1.4×10−31.4× 10^-3 4.0×10−34.0× 10^-3 Rank r 256 512 Basis update interval K 10 (β1,β2,β3)( _1, _2, _3) (0.95,0.95,0.95)(0.95,0.95,0.95) Compensation weight αc _c 1.0 Norm-growth limit γ 1.01 Learning rate 4.4×10−34.4× 10^-3 1.0×10−31.0× 10^-3 (β1,β2)( _1, _2) (0.95,0.95)(0.95,0.95) Rank r 256 512 Oversampling p 8 16 Power iterations 1 Preconditioning frequency T 10 Shampoo β (per refresh) 0.95 Max. precond. dimension 4096 8192 Clean (Ours) Norm-growth limit γ 1.01 Learning rate 4.4×10−34.4× 10^-3 1.0×10−31.0× 10^-3 Other hyperparameters as Clean Eigenbases UL,URU_L,U_R INT8 linear, group 256, SR First moments INT8 linear, group 256, SR Second moments INT8 dynamic, block 256 Q-Clean (Ours) Eigenvalues as Clean