Towards Efficient Optimizer Design for LLM via Structured Fisher Approximation with a Low-Rank Extension
Authors: Wenbo Gong
Meyer Scetbon
Chao Ma
Edward Meeds
Venue: ICLR 2026 Poster (2026)
Keywords: Optimization
LLM training
memory efficiency
Primary area: optimization
TL;DR: We unify many recently introduced LLM optimizers under approximating Fisher information with different structural assumptions. Based on this we propose two new memory-efficient optimizers with 2x faster convergence than Adam.

Abstract: Designing efficient optimizers for large language models (LLMs) with low-memory requirements and fast convergence is an important and challenging problem. This paper makes a step towards the systematic design of such optimizers through the lens of structured Fisher information matrix (FIM) approximation. We show that many state-of-the-art efficient optimizers can be viewed as solutions to FIM approximation (under the Frobenius norm) with specific structural assumptions. Building on these insights, we propose two design recommendations of practical efficient optimizers for LLMs, involving the careful selection of structural assumptions to balance generality and efficiency, and enhancing memory efficiency of optimizers with general structures through a novel low-rank extension framework. We demonstrate how to use each design approach by deriving new memory-efficient optimizers: Row and Column Scaled SGD (RACS) and Adaptive low-dimensional subspace estimation (Alice). Experiments on LLaMA pre-training (up to 1B parameters) validate the effectiveness, showing faster and better convergence than existing memory-efficient baselines and Adam with little memory overhead. Notably, Alice achieves better than 2x faster convergence over Adam, while RACS delivers strong performance on the 1B model with SGD-like memory.

Published as a conference paper at ICLR 2026
TOWARDS EFFICIENT OPTIMIZER DESIGN FOR LLM VIA STRUCTURED FISHER APPROXIMATION WITH A LOW-RANK EXTENSION
Wenbo Gong" Microsoft Research
Meyer Scetbon"1
Citadel Securities
wenbogonggmicrosoocoo com
meyerscetbon gmail com com
Chao Ma¹ Microsoft Research
Ted Meeds
Microsoft Research
waw..mmmmmiccooooooooooo com
Meeds cmicrooft.ocooooooocoooooo
ABSTRACT
Designing efficient optimizers for large language models (LLMs) with low-memory requirements and fast convergence is an important and challenging problem. This paper makes step towards the systematic design of such optimizers through the lens structured many state-of-the-at efficient optimizers can be viewed as solutions to approx imation (under the Frobenius norm) with specific structural assumptions. Building on these insights, we propose two design recommendations of practical efficient optimizers for LLMs, involving the careful selection of structural assumptions to balance generality and efficiency, and enhancing memory efficiency of optimiz- ers with general structures through a novel low-rank extension framework. We demonstrate how to use each design approach by deriving new memory-efficient optimizers: Row and Column Scaled SGD (RACS) and Adaptive low-dimensional subspace estimation (Alice) Experiments on LLaMA pre-training (up to pa- rameters) validate the effectiveness, showing faster and better convergence than existing memory- memory-efficient baselines and Adam with little memory overhead. No- tably. achieves better than 2x faster convergence over Adam, while RACS delivers strong performance on the 1B model with SGD-like memory.
information matrix approximatioiiooooo We show that
INTRODUCTION
Adaptive optimizers are critical training large language models (LLMs). Yet, as models and datasets continue grow, it brings several implications for training, including increased GPU requirements reduced per-device batch size, which lowers overall training throughput. Adam, for instance, triples memory requirements due the storage of internal eexoneeniil moving average (EMA) states, while other optimizers (Gupta et al., 2018 Vyas et al. (2024) with faster convergence (in terms of training steps) can further inflate the total memory. Meanwhile. some memory efficient optimizers, like stochastc gradient descent (SGD) fail to train the LLMS effectively. Thus, designing efficient optimizers has become increasingly important.
There are several promising lines of research that have advanced one or more of these aspects of efficiency: (1) removing the optimizer state (Ma et al.) 2024 Jordan al.) (2024] Zhang al.) | [024]] Xu et al. 2024 Zhu et al. 2024), or (2) using ad-hoc low-rank design to reduce memory (Hu et al.
2021 Lialin et al. 2023) Zhao et al. 2024a Chen al. 2024a| et al. 2024). However, developing
new efficient optimizers remains challenging. This paper explores structured FIM approximation potential framework for this task.
To demonstrate the effectiveness of this framework, we begin by showing that many existing opti- mizers and gradient operators, including Adam, Shampoo, gradient normalization and whitening
Equaa contribution
1 This work is done when Meyer was affiliated with Microsoft Research

---

Published as a conference paper at ICLR 2026
(Zhang et al. 2024 Gupta et al., 2018 Vyas et al. 2024; You et al. 2019 Ma et al. 2024 Jordan et al. 2024). can be recast under structured FIM approximation. We then go on to show how to derive two generalizations of Adam with more relaxed structural assumptions, named One-sided SOAP (OSOAP) and SOAP/AdaDiag++ (Vyas et al., 2024 Nguyen et al. (025). Previously, (OSOAA (OSOA] is considered simple one-side projection version of SOAP, but we reveal that they correspond to different underlying approximation structures. Although this framework provides a clear link between structures and optimizers, working with more general structures can come at the cost of efficiency. For example, SOAP can require 7 times more memory than SGD.
Building on these insights, our first design recommendation proposes to choose structural assumptions that balance generality with practical efficiency. We demonstrate this by choosing structure that generalizes gradient normalization, leading to a new efficient optimizer, Row and Column Scaled SGD (RACS) with SGD-like memory requirements.
For optimizers with more general structures it may not be always possible to achieve such a balance. Instead of rejecting their potentials, our second design recommendation proposes to apply novel low-rank extension framework to improve their efficiency. This framework consists of three steps that can convert full-rank optimizers into their low-rank approximations with reduced memory and computational costs. We demonstrate this by deriving a low-rank extension of [OSOAP] called Adaptive don-dimensional subspace estimation (Alice) We also prove its convergence under the continuous-time setup.
We emphasize that the main contribution is the unifying viewpoint that allows systematic design for efficient optimizers. RACS and Allie are two examples. More novel optimizers can be derived by (I) designing new structural approximations; or (2) adapting existing optimizers with our low-rank extension. This allows one to go beyond the ad hoc design and focus on well grounded strategy.
With pre-training LLaMA Touvron al. 2023) with C4 dataset (Raffel et al. 2020), demonstrate that RACS and | outperform Adam and several memory- efffccient baselines. [Alice achieves better than 2x speed-ups compared Adam, and RACS performs strongly on pre-training the IB LLaMA.
To summarize, our contributions are:
We propose structured FruuttuedFFMMapppreee approxxmaaii practical framework for optimizer design and show that several existing optimizers can be viewed as special cases within this framework.
We propose design optimizers by choosing structures that balance generality and efficiency and demonstrate it with a new optimizer, RACS
We present a low-rank extension framework to convert full-rank optimizers and demonstrate it with a new optimizer Alice We also prove its convergence property.
We demonstrate the effectiveness of RACS and Alice on LLaMa pre-treeeee pre-training tasks when compared to Adam and other baselines.
summarize the main relationships among optimizers and the procedure of our proposed low-rank extension framework.
Fig. and Fig. 2
Scope of this paper We want to emphasize and clearly define the scope of this paper. Specifically, we DO NOT aim to provide the definitive and complete answers on how to design an optimizer for
LLM and their theoretical properties. Instead, our primary contribution is to establish that structured
FIM approximation is a potential candidate as a unifying design principle for LLM optimizers, followed by two design guidelines with exampled efficient optimizers. In summary, we aim to provide a potential direction for future LLM optimizer research.
'Generality of structural assumption defines how relaxed this assumption is. We say structure A is more general than B iff. B can be recovered hy applying further constraints on A. The general structures tend to give hetter approximations Fo the less general one.
2

---

Published as a conference paper at ICLR 2026
Haaidy
Low wwA.SSITTETIII
LoLow-lann Trecong
Swiscing
Comperadan
Figure 2: Summary of the low-rank extension framework to convert full-rank optimizer into its low-rank counterpart.
Figure 1: Summary of relationship between differ- ent optimizers and their corresponding structural assumptions.
2 PRELIMINARIES
2.1 BASIC NOTATIONS AND SETUP
Throughout the paper we consider 2D matrix parameters W (i.e. layer weights) of size x n and m & nı the operation Vec(.) vectorizes the input matrix by stacking its columns; Mat(-) is the inverse of Vec(.) and reshapes the vector back into matrix. We use Lo as the loss function where = Vec(W). G = w L is the matrix gradient and g is the vectorized gradient, i.e. g = Vec G). gi denotes the 7th column of G. In the paper, we assume W are parameters of one layer. & indicates the Kronecker product. By default, M and VM indicates the elementwise multiplication and square-root. EVD(M,r) r) performs the eigen-value decomposition (EVD) and keeps the top r eigenvectors ordered by the descending eigenvalues. If is omitted, we keep all eigenvectors. QR(M) with orthonormal M E Rm obtains the remaining m r basis via QR decomposition.
We also introduce the following diagonal operations: (i) will extract the diagonals of M to vector. Diags [M1... M,) will stack the input sequence of matrices to form a larger block diagonal matrix. Diag. (v) will expand the imput vector v to a pure diagonal matrix. Diagm (M) will expand the input M to a larger pure diagonal matrix by stacking the element of M in column-wise order. We demonstrate the above operations with examples Sec. C
Note that results in this paper may involve the expectation E[-]. we assume the use of EMA as its practical estimation.
2.2 FISHER INFORMATION AND NATURAL GRADIENT DESCENT
Fisher information a fundamental concept in statistics that measures the amount of information a random variable carries about parameter of interest. In this paper, we ignore dependence between layers and treat each layer independently. Under the context of LLMs, with the vectorized mini- batched gradient of one layer g. we define the empirical FIM for that layer as F Elgg"). Natural gradient descent (NGD) leverages the inverse of FIM to smooth the local geometry to improve convergence and stabilize training (Martens) 2020 In practice, the square-root inverse FIM may
be preferred due to its properties and improved performance on certain optimization problems Yang
& Laaksonen 2008 Lin al.) [2024] Loshchilov & Hutter 2016: Bergstra & Bengio 2012 Choi
2019). The corresponding update of W is:
W - W - A AM
(I)
Due to the large size of F [ computing its inverse is a significant impediment to applying this update rule in practice. One way address this is to enforce certain structures to approximate FIM In this paper, we consider two main classes of structures: block diagonal and Kronecker product matrices. They possess favorable properties that can significantly reduce computational complexity. Sec. D.2 provides a brief introduction and summary of their key properties. We also include a comprehensive background in Sec. D be more self-contained.
3

---

Published as a conference paper at ICLR 2026
Table I: The summary of underlying structure assumptions of existing and newly proposed optimizers, along with practical efficiency. "Generalizes" refers to the optimizer whose structure is generalized, e.g. OSOAP]s structure generalizes Adam (see Sec. for further details). Optimizers blue color are new. The "Computation" is the cost of per-step of the optimizer update. We assume m >
Adam
Stasipoo
Miniig
Gnllore
RACS
Alic
Stoctue
Dog.(o)
Rİ
Disgnt (UDu((U-D.UU)))
Un UnDDUUn Ge)
Approx.
$eQ
Godam
Geneniaas
NIA
NIA
Adam
Shampoo
NGA
Nonmlication
Componation
Otson)
Ojas
O(m?)
Olyoor maa/K)
O(son) von
Qimer =xPx/
Menoyy
mo
update
STRUCTURAL APPROXIMATION TO FIM
The structured FIM approximation framework consists of three steps: first, choose structure to approximate FIM second, find the solution F by minimizing the following objective:
min FER
| FlI::
where F is the empirical FIM H is the matrix family corresponding the structural assumption; third, derive the square-root NGD update with approximated F. In this section, show that many existing optimizers and gradient operators can be reformulated within this framework by varying structural assumptions, thereby establishing connections between the structure and optimizers.
3.1
ADAM: PURELY DIAGONAL STRUCTURE
When Adam was originally proposed. was argued that the purpose of the second moment was to approximate the diagonal elements FIM (Kingma 2014, Hwang 2024).
Proposition 3.1 (diagonal approximation). Assuming H (Diag, (v); > 0}, then Eq. (2) has analytic solution F = Diag, (B(g²)) where g² indicates the dlement-wise square of g = Vec(G)).
is trivial to show the resulting square-root NGD recovers Adam's second moment when using the EMA to estimate E. Together with the first moment, one can recover Adam.
3.2
SHAMPOO: KRONECKER PRODUCT STRUCTURE
Although previous work (Morwani al. 2024) has demonstrated the connection of Shampoo (Gupta et al. 2018) to the Kronecker product approximation through power iteration algorithm, we provide an alternative view: Shampoo's pre-conditioner can be derived by minimizing an upper bound of Eq. ( with
H [RA
Li,; R.
R Lm E 1mXm)
where R. and L,, are symmetric positive definite (SPD) matrices.
Theorem 3.2 (Shampoo pre-conditioner). Assume the above structural assumption, then we have an upper bound of Eg. (2)
F - 3
3 mm 2 12 B² - C + mm CTP -
[+IF++1B² (²I
where A R Im B I, L C E Elgg"1 b and? 2 indicates the matrix square. Minimizing the upper-bound admits Rt = = LE G] and L;, = LEGG
Sec. F. shows that the corresponding square-root NGD leads to the Shampoo's 's update formula. Therefore, the structure behind Shampoo is the Kronecker product two square-root SPD matrices.
3.3 NORMALIZATION AND WHITENING: SIMPLE BLOCK DIAGONAL STRUCTURES
For an input G, the whitening and normalization operator defined
Whitening(CC G) = ( Norm(G) = GS-
where Diag(S) contains the squared column 12 norm of G. Next, we show that these two operators also corresponds to minimizing Eq. with different structural assumption.
4

---

Published as a conference paper at ICLR 2026
Proposition 3.3 (Normalization and whitening). With H (I, & M and H. {s & Imi S O, Vi where M E Rmxx is SPD and E Ru** is a positive diagonal matrix. Then, minimizing Eq. admits
M E|GG"1
(3)
n
S" = E Diag. (gf 91. .t.......................... ........hh..))
gm))
(4)
The proof can be found in Sec. G.3 where we prove a much more general solution (Theorem G.4).
and Theorem 3.3 is special case. The square-root NGD update with Eq. (3) and Eq. (4) are
Mat (F--F---- 1 gg) VnB[GG") and /mGGS*= respectively (refer to Sec. F.2 for derivation).
We can view Whitening()) and Norm(.) as special cases with one-sample estimate for E.
Muon and many others are special cases Many recently proposed optimizers, such as Muon, SWAN, LARS, and LAMB (Jordan al., 2024, Ma al. 20244) You et al.) [017| 2019), rely on normalization and/or whitening. These gradient operators improve convergence You et al. 2017
2019 Jordan al. 2024) and can remove Adam's internal states (Ma et al.) 2024). In fact, Muon simply whitens the gradient, and thus is special case of this framework. Interestingly, the original Muon lacks vn coefficient predicted by Theorem 3.3 This was later added back by (Liu et al. 2025). which is proved to be critical for LLM training. We include thorough discussion Sec. 1.5
3.4 OSOAP AND SOAP: TWO GENERALIZATIONS OF ADAM
All structures considered above are do not strictly generalize Adam's 's purely diagonal structure. Next, consider two structures that strictly generalize Adam, normalization, and whitening: block diagonal matrix with a shared eigen-space; and its combination with the Kronecker product.
First, consider block-diagone structure that generalizes Adam:
(5) where y defines shared full-rank eigen-space, and D, is eigenvalue matrix. Adam is special case by constraining , I. The structures in See. 3.3 are also special cases by setting D; = D (whitening); and U I, D; s; (normalization). However, this structure does not lead to analytic solution for Eq. (2). Instead, we propose an approximation by solving a 1-step refinement.
H. = (Diage (M, [M,..., M,); M, = U,D,U7)
Theorem 3.4 (1-step refinement). For the structure in Eg. (5). we consider the following 1-step refinement: (i) assuming all D; are equal, and find Uf; (ūi) fix the U, and find D;. Then, step (i) admits the analytic solution from Theorem 3.3
U EVD(EIGG"))
(6)
where EVD is the eigenvalue decomposition; step (ii admits the analytic solution:
D =Diag (E[ G)
(7)
where D Diage (Di..., D,).
Based on this result, we can derive the corresponding square-root NGD referr to Sec. F33
U
Mat F F- UU,
(8)
B
This can be viewed as applying Adam's update on space "rotated" by Uy. Consequently. we can design an optimizer based on Eq. (8), calledOSOOi The following procedures describes its update:
= 3 3,mm-1 + (1 G.
Q 13 + (1 G. U = EVD(Q.)):
A - U.
1i = -3,00-11 + (1 -1- (1-32) (UT,))
Algorithm 7 in Sec. D.6 summarizes its algorithm. In Vyas et al. 2024). they regard OSOAP and SOAP as being from the same optimizer family and the former as the one-sided projection verion of the latter. In the following, will show SOAP and OSOAP correspond to different underlying structure assumptions, indicating that (S0APishoule not simply treated as the cheaper approximation to SOAP.
5

---

Published as a conference paper at ICLR 2026
SOAP/AdaDiag++ Despite the structural generality of OSOAP it only focuses on the block diagonals of FLM To go beyond this, one can combine the structure of with Shampoo's s Kronecker product to obtain structure that generalizes all the others in this paper. Interestingly, this turns out to be SOAP/AdaDiag+++ Refer to Sec. for details.
4 RACS| MEMORY-EFRICIEN OPTIMIZER FROM CAREFULLY SELECTED STRUCTURE
The structured FIM approximation reveals two important insights: there exists a correspondence between structural assumption and optimizers, and structural generality often comes at the cost of practical efficiency. For example, while the structures OSOAP and SOAP offer more accurate McMM approximations, they require expensive computation and memory consumption (Table D Building on this, our first design recommendation is to select structures that balance generality and efficiency.
To demonstrate this, we select structure that generalizes gradient normalization, which scales both the rows and columns simultancously:
H = 1 @ )}
(10)
where S E Ruxx, Q E Ruxm are positive diagonal matrices. The optimal solution of Eq. can be solved by fixed-point procedure:
Proposition 4.1 (Two-sided scaling). Assuming the structure of Eq. 100 and E[G
contains only
positive values, solving Eq. (2 admits an iterative fixed point procedure:
Diag (E|G"QQT
Diag (EIGSG")
1Q1
q
||
(11)
where s Diag( S), q = Diag(() Additionally, the fixed point solution s, q converges the right and left principal singular vector of E up to scaling factor with unique S @ Q.
The corresponding square-root NGD update scales both rows and columns through S and Q (i.e. Mat(F- = Q GS- We name this optimizer, Row and Column Scaled SGD (RACS) (Algorithm in Sec. BB. Although analytic solutions of s. q exist, we perform 5 steps of Eq. (1] as an efficient approximation. To further stabilize training, we also incorporate the norm-growth limiter used (Chen et al. 2024a), RACS is highly memory efficient since only needs the storage two diagonal matrices S and Q and a scalar for the limiter, consuming m + n + memory. In Sec. [.5 we discuss connections to Adapprox, Apollo and Adafactor (Zhao et al. 2024b Zhu et al. 2024 Shazeer & Stern 2018). The above structure may not always be casy to find. Therefore, our next design guideline based on the novel low-rank extension is more widely applicable.
5 ALICE MEMORY-EFFICIEN OPTIMIZER FROM LOW-RANK EXTENSION FRAMEWORK
We propose a novel low-rank framework consisting three steps, low-rank tracking, subspace switching; and compensation. Tracking aims to reduce the memory cost oMttssstees,,, whereas switching and compensation are designed to correct the potential issues caused by low-rank approx- imations. We demonstrate these by converting OSOAP to its low-rank version, Alice While the procedure could be applied to SOAP in similar manner, we leave this for future work.
5.1 TRACKING: LOW-RANK PROJECTIONS TO REDUCE MEMORY
By carefully examining the OSOAP procedure (Eq. ()),, ). we notice the key internal state of OSOAP the projection U, To reduce the memory cost, we propose low-rank U, by keeping only the top eigenvectors, and denote the remaining m r basis as Uc., (i.e. U [U., Uc,)). For Q, we propose the following low-rank tracking:
a, -.U Qun = + Ba)a
(12)

---

Published as a conference paper at ICLR 2026
where o, is the projected gradient, and Qi is the low-rank tracking state. One can easily reconstruct back Q: U when needed. This reduces the memory from m² of Q, r2 of Qt-
However, this low-rank projection comes with two consequences: (1) the reconstructed state Q2 is no longer accurate; (2) the resulting parameter update A in Eq. (9) ignores the information within Ue.t- Next, we propose two additional steps, switching and compensation, rooted in theoretical insights to address the two problems, respectively.
5.2 SWITCHING: MIXING LEADING BASIS WITH THE COMPLEMENTS
Since the projected gradient or only maintains the information in the space spanned by U, the low- rank tracking state Q: necessarily discards information in Ue,t- Therefore, even if those directions should become the leading basis, Q2 will ignore their contributions, causing the stability of the previous leading cigenvectors and preventing the exploration of other spaces. This effect is largely ignored by the previous work (Yen et al. 2023 Zhao et al.) 2024a) and believe one should always keep the top eigenvectors. Here, challenge this commonly adopted concept.
Theorem G.8 confirms this and shows that the true tracking state Qt can be decomposed into a low-rank reconstructed state (i.e. U,Q,U77 and a residual term quantifying the importance of Ue,t- When the residual term becomes dominant, the remaining basis will become the new leading basis. Inspired by this theoretical insight, we propose a practical scheme to replace the original EVD in Eq. (8), called subspace switching (Algorithm []. deliberately replace some leading eigenvectors with randomly sampled basis, forcing the optimizer to explore other spaces.
Algorithm Subspace switching
Algorithm Compensation
1: Input: Reconstructed state Q, rank r, lead- ing basis number l, previous low-rank projec-
1: Input: G, projection U, previous norm
limiter norm Φ, threshold 7, decay rate 3
tion Ut-1
2: p Bp + (1 B) G ["
2: U Subspace iteration( Q, U,-1)
Compute C, as Eq. (4]
3: U + Keep top eigenvectors of U;
4: n = 7/ max ( ||C,
7 if Φ > 0 else
4: Uniformly sample r basis from the com-
5: & |||/C;||
plement U², QR(U)) U222
C. t n/C
5: Combine basis U [U/,1 0221
7: Return C, P.
6: Return U
5.3 COMPENSATION: CONVERT LOW-RANK UPDATE TO BE FULL-RANK
Another problem with low-rank projection U, is the information loss in the resulting parameter update A at each step. The goal of compensation step to compensate for this information loss with minimal memory overhead. We first show that the full-rank update for A with U, in Eq. (8) can be decomposed into the low-rank and complement update controlled by U, and U (proof Sec. [.6]. We omit the subscript t for simplicity:
UTG
UTG
A =
U.
(13)
VE[U*G
VE
Mat(F, F. 1/0)
is the approximated FIM corresponding to the where F. = Diage(UU....U.) ....., U.Denn U?) "..DD
complement basis, D.,, Diag, (B((UT) (I) and -3 is the square-root pseudo-inverse.
We notice that Mat 2 g) is exactly the square-root NGD with FIM F We can directly leverage the FIM view point to approximate F. with a carefully selected structure that leads to memory-efficiient compensation procedures. One choice is the diagonal structure of normalization operator. Based on this, propose the following compensation at each step:
C Mat((S U.U. )g) U.U... GS
(14)
7

---

Published as a conference paper at ICLR 2026
where S is positive diagonal matrix, U. = (G - GU is the gradient information
preserved in the remaining basis. This compensation admits an analytic solution S to Eq. ( with Fe as the target (Theorem G.9 in Sec. G.
m
Diag(S) V 1
(15)
Algorithm 2 summarizes the compensation step with the norm growth limiter from (Chen et al. 2024ap. C. will then be used to replace the Mat (F2 g) Eq. (3) for each update step.
5.4
ALICE OPTIMIZER AND ITS CONVERGENCE
Based on the above analyses, we propose Alice and a variant Alice without tracking (Alice-0) as low-rank extensions to OSOAP Interestingly. GaLore, in fact, is an approximation to Alice when disabling tracking, switching and compensation. Our view point offers connection of GaLore to OSOAP revealing that it is a low-rank version of more powerful optimizer than Adam, and providing an explanation of its effectiveness. Algorithm 4 in Sec. summarizes Alice procedures.
To further support the theoretical validity Alice we prove its convergence under the continuous- setup (Maddison et al., [2018] Nguyen et al. 2005) Chen et al.] 2023) We
Convergence of Alice
leave the general treatment of the convergence to future work. Please refer to Theorem H.2 in Sec. H
Connections to structured FIM approximation The proposed low-rank framework seems to be separated and independent to the structured FIM viewpoint. On the contrary, they are deeply connected. First, |iIMaaaprxximmaiion provides a base optimizer as a starting point for the low-rank extension to work on. More importantly. the critical compensation step is entirely inspired the structured FIM viewpoint. Since the information loss from low-rank projection has NGD format (Eq. (13) with a FIM like preconditioner (Fa). we a principled question: How can we compensate this in an optimal way? Through structured FIM viewpoint, we aim to approximate this F, with simple structural assumptions, leading to compensation term with minimal resources requirement.
EXPERIMENTS
6.1 PRETRAINING LLAMA WITH C4 DATASET
Setup We evaluate the proposed RACS Alice and its variant Alice-O pre-training LLaMA (Touvron et al., (023) on the C4 dataset (Raffel et al. 2020). We train the following model sizes: 60M, 130M, 350M and 1.3B using similar setup to (Zhao al. 2024a Zhu et al.] 2024). An important consideration in our experiments is that all previous low-rank methods rely on full-rank Adam to train the last layer, which is arguably one of the most important layers (Zhao al.) 2024c). To thoroughly assess their effectiveness, we report performance for both cases for low-rank methods- -traissngnnnngggniiii the last layer with and without Adam but prioritize the latter. For baselines, we consider GaLore, Fira, Apollo-mini Apollo-svd, Muon and Adam. The priority of this paper focuses on the memory-efficiency of optimizers. Therefore, we mainly adopt the training setup from Zhao et al. (2024a), where all training are under BF16 (i.e. model, parameter state, etc.). We leave the exploration of mixed-precision training in future work.
Main results Table 2 reports the pretraining performance in terms of evaluation perplexity, RACS and Alice outperffrms the other memory memory-efficcint baselines and full-rank Adam consistently across model sizes. and Allice-0 on par with each other, suggesting (Alice-0 may be preferred for better memory efficiency. One major advantage of Alice is its fast convergence compared to Adam, and achieves more than 2x speed-ups across model sizes, while maintaining similar memory consumption as GaLore. Fig. shows the eval ppl. for the IB model. Despite the simplicity of RACS, it performs well for 350M and 1.3B model, and surpasses the other scaling- based baseline, Apollo-mini and Apollo-svd, consistently. From the throughput (TP), we observe Alice and RACS is not significantly slower than Adam with 15% and 11% drop for IB model. Considering the speedup effect, report the effective TP represent how quickly the optimizers reach target loss
nuggesting/Alice-[]]
8

---

Published as a conference paper at ICLR 2026
Table 2: I.LaMA pretraining performance. Ppl. is the evaluation perplexity, Mem. is the estimated memory consumption of optimizers, Regarding the last layer, we report ppl of GaLore and Fira for both cases, For baselines, the ppl. in bracket are directly cited from (Zhu et al. 2024). where the last layer is trained by Adam and consumes more memory. We also cite the numbers of Apollo-mini, Apollo-svd from (Zhu et al.) 2024); and Muon from (Ma et al,, (0044). The speed-up is measured against Adam in terms of steps. The TP is the number of training tokens processed per second, and
effective TP is total training token used by Adam divided by time used by the candidate optimizers to reach the final eval ppl of Adam. Effective TP is equivalent to wall-clock time performance.
Methods
60M
130M
350M
1.00
Ppl
Mem
Pyt
Man.
Pyl
Man.
Fni
Man.
Adam
33.94
0,320
25.03
0,75G
19:34
2,053
16.44
7,48G
Ool
38 3891154555
0.21 2033
27.18 (25 16)
0.516005701
21 (13
(30(1290)
4250444480)
Firs
33.77
0.21GH0.20 2603)
25.21
(22.75)
0.11G00570)
(893(17.09)
1,30 29()
15. (14.4(11.11)
4.20((4)4)))
Apall--minii
31,93
0,25G
23.84
043G
17.18
0.33
14:17
2.96G
Apollossv
31,36
0.266
22.84
0.57G
16.07
1.29
14:10
4,43G
Muon
31.6
0.27G
24.59
0
19.30
1.493
16.08
5,23G
RACS
10.25
0.23
32.67
0.436
16.51
0.93
13.52
2.56G
28.83 74)
0.21 263)
21,99 (22.43)
0.51G
16,66 (16.43)
1.20 29(3)
13,97
(1347))
4,500 (4.44()
Alice
28.69
0,22Gr0.263)
21.95
(31.79)
0.53G
16.11 (16.37)
(1,36)
13,85
(13.52)
(4.6G)
Speed-up (Alice)
2,22X
2000
2,45
2,82x
Thoughgu PPyteeee Wd)
9174899774
82249722447
631239631399
53588/53585
Thoogggus: Pyytlte
92589/202658
71985144145
58847/143088
45523123048
Throghous: PPPitee
98228/624225
73253/123116
55970/131372
47485/129817
Eval PRL with model siee: 1b
Method
Mem.
40K
80K
120K
1.50K
22
8-bir Adam (78)
266
18.19
13.47
14.83
14.61
Asami
8-bit Galore (7B)
18G
17.94
15.39
14.95
14.65
Apollo (7B)
15.03G
17.55
14.39
13.23
13.00
RKCS
Apello-mini (7B)
13.53G
18.03
14.60
13.32
13.09
17
Colore
Fiu
(IB)
2.98G
14.26
13,20
XTlie (R)
4.60
15.93
14.08
13.15
12.97
Iow 6-lmn
1.2
16.600
Table 3: Eval ppl. at different training steps. For baselines, we cite the number from (Zhu al. 2024).
12
500
15.67
Components
Eval ppl.
No tracking, switch, compen. Tracking
26.96
12
36000
(0003
Update Step
60000
90003
100003
27.35
Tracking+Switch
25.11
Figure 3: 1B LLaMAA eval ppl. curve. "+Im head" represents the last layer trained by Adam.
Tracking-Swicch-Coonpe
21.95
Table 4: Effectiveness of each components.
in wall-clock time. Alice and RACS achieve 123048 and 129817, respectively, compared to 53588 for Adam, resulting in more than 2x faster convergence in wall-clock time to reach Adam's 's final evaluation perplexity. For the actual memory footprints and additional training curves, see Sec. K.5
1B v.s. 7B LLaMA To further demonstrate the effectiveness, we follow the setup of (Zuu et al. 2024), but train a IB Alice and RACS to compare with 7B LLaMA trained by baselines. Table shows that IB model with optimizer consistently outperforms 7B model trained with baselines at different training checkpoints. For 150K steps with 100, Apollo requires 15 days while Alice requires 3.8 days. We DO NOT claim 1B can outperform 7B model trained with standard AdamW, but demonstrate the effectiveness of our optimizers compared to other memory-efficient baselines.
6.2 ABLATION AND GLUE FINETUNING
Effect of low-rank tracking and switching strategy First, we verify whether low-rank tracking (Eq. (2)) is beneficial. As shown in Table 2 Alice-O performs par VitheAlice, suggesting that tracking does not provide a significant boost. Fig. 8(a) indicates that tracking can be helpful when
compensation is disabled and used alongside switching. Without switching, tracking leads to inferior performance due to the stability of the eigenspace (Fig. We also evaluate the effectiveness of switching strategy compared to other heuristics: (1) Gaussian: U is sampled from Gaussian distribution: (2) Gaussian mix: the leading basis is mixed with basis sampled from a Gaussian matrix; (3) full basis: instead of sampling the basis in Algorithm [ solely from m complement

---

Published as a conference paper at ICLR 2026
basis, it is sampled jointly from the top I and basis. Fig. outperforms the alternatives with clear margins.
8(b) shows that our strategy
Compensation strategy, other ablations and GLUE finetune The closest work our compen- sation step is Fira (Chen et al.] 2024a), which introduces a heuristic-based compensation term. To evaluate the effectiveness of our optimal compensation, we compare it against Fira compensation and no-compensation baseline. As shown in Fig. [[c) our optimal compensation achieves better performance and convergence than all candidates. Table 4 summarizzs the contributions of each component. We also performed additional ablations: (1) the effect of the last layer, (2) the effect of [MA in RACS] and (3) hyperparameter sensitivity analysis, and GLUE (Wang) 2018) finetuning experiment. For more details, see Sec. K.6land Sec. K.7 for details.
7 CONCLUSION AND LIMITATIONS
In this paper, we take step toward the systematic design of efficient optimizers for LLMs through structured FIM approximation. We establish the connection between structural assumptions and optimizers. Building this, propose design approaches: (1) simple structure that balances generality and efficiency; (2) low-rank extension framework for more general structures. These are demonstrated with novel optimizers, RACS and Alice, respectively.
Despite the theoretical and empirical results, this paper does not aim to provide a complete recipe but serves starting point for future research. For example, complete convergence treatment for all optimizers under this FIM framework is an important research question that is currently lacking in this paper. Also, how incorporate the model structure and task into the optimizer design is another promising direction.
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer. Scalable second order optimization for deep learning. arXiv preprint arXiv:2002.09013. 2020.
James Bergstra and Yoshua Bengio. Random search for hyper-parameter optimization. Journal of machine learning research, 13(2), 2012.
Jeremy Bernstein, Yu-Xiang Wang, Kamyar Azizzadenesheli, and Animashree Anandkumar. signsgd: Compressed optimisation for non-convex problems. In International Conference on Machine Learning, pp. 560-569. PMLR, 2018.
Lizhang Chen, Bo Liu, Kaizhao Liang, and Qiang Liu. Lion secretly solves constrained optimization: As Iyapunov predicts. arXiv preprint arXiv:2310.0589) 2023.
Xi Chen, Kaituo Feng, Changsheng Li, Xunhao Lai, Xiangyu Yue, Ye Yuan, and Guoren Wang. Fira: Can we achieve full-rank training of Ilms under low-rank constraint? arXiv preprint arXiv:2410001623, 2024a.
Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang, Hieu Pham, Xuanyi Dong. Thang Luong, Cho-Jui Hsieh, Yifeng Lu, et al. Symbolic discovery of optimization algorithms. Advances in neural information processing systems, 36, 2024b.
Wonwoong Cho, Sungha Choi, David Keetae Park, Inkyu Shin, and Jaegul Choo. Image-to-image translation via group-wise deep whitening sebeeninn-and--oor and-coloring transformation. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pp. 10639-10647,, 2019.
D Choi. On empirical comparisons of optimizers for deep learning. arXiv preprint arXiv: 1910.5491. 2019.
Sungha Choi, Sanghun Jung, Huiwon Yun, Joanne T Kim, Seungryong Kim, and Jaegul Choo. Robustnet: Improving domain generalization in urban-scene segmentation via instance selec- tive whitening. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 1 11580-11590, 2021.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 2011.
Kaixin Gao, Xiaolei Liu, Zhenghai Huang, Min Wang, Zidong Wang, Dachuan Xu, and Fan Yu. A trace-restricted cronecker-factoee approximation natural gradient. Proceedings the AAAI Conference on Artificial Intelligence, volume 35, pp. 7519-7527, 2021.
Thomas George, César Laurent, Xavier Bouthillier, Nicolas Ballas, and Pascal Vincent. Fast approximate natural gradient descent in kronecker factored eigenbasis. Advances in Neural Information Processing Systems. 31, 2018.
Roger Grosse and James Martens. A kronecker-factored approximate fisher matrix for convolution layers. In International Conference on Machine Learning, pp. 573-582. PMLR, 2016.
Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimiza- tion. In International Conference on Machine Learning, pp. 1842-1850, PMLR, 2018.
Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang. Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. arXiv preprint arXiv:2106.09685 2021.
Lei Huang, Yi Zhou, Fan Zhu, Li Liu, and Ling Shao. Iterative normalization: Beyond standardization towards efficient whitening. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 4874-4883 2019.
Dongseong Hwang. Fadam: Adam a natural gradient optimizer using diagonal empirical fisher information. arXiv preprint arXiv:2405. 12807, 2024.

---

Published as a conference paper at ICLR 2026
Keller Jordan, Yuchen Jin, Vlado Boza, Jiacheng You, Franz Cecista, Laker Newhouse, and Jeremy Bernstein. Muon: An optimizer for hidden layers in neural networks, 2024. URL https
//kellezjordan. http:aam....ttti...oootottts.... http:/ossppss/mmooooooooooo
Jean Kaddour, Oscar Key, Piotr Nawrot, Pasquale Minervini, and Matt J Kusner. No train no gain: Revisiting efficient training algorithms for transformer based language models. Advances in Neural Information Processing Systems, 06.25793-25818, 2023.
Diederik Kingma. Adam: method for stochastic optimization. arXiv preprint arXiv; 442.66880. 2014.
Abdoulaye Koroko, Ani and Quang Huy Tran. Efficient approximations of the fisher matrix in neural networks using kronecker product singular value decomposition. arXiv preprint arXiv:2201.10285, 2022.
Ibtihel Ben Gharbia, Valérie Garès, Mounir Haddou,
Peihua Li, Jiangtao Xie, Qilong Wang, and Zilin Gao. Towards faster training of global covariance pooling networks by iterative matrix square root normalization. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 947-955, 2018.
Xi-Lin Li. Preconditioned stochastic gradient descent. IEEE transactions on neural networks and learning systems, 1955):1454 1466, 2017.
Yijun Li, Chen Fang, Jimei Yang, Zhaowen Wang, Xin Lu, and Ming-HHsund Yang. Universal style transfer via feature transforms. Advances in neural information processing systems, 30, 2017.
Vladislav Lialin, Sherin Muckatira, Namrata Shivagunde, and Anna Rumshisky. Relora: High- rank training through low-rank updates. In The Twelfth International Conference on Learning Representations, 2023.
Wu Lin, Felix Dangel, Runa Eschenhagen, Juhan Bae, Richard E Turner, and Alireza Makhzani. Can we remove the square-root in adaptive gradient methods? second-order perspective. arXiv preprint MtXii:240220334966 2024.
Hong Liu, Zhiyuan Li, David Hall, Percy Liang, and Tengyu Ma. Sophia: A scalable stochastic second-order optimizer for language model pre-training, arXiv preprint xxxx:22305.14342 2023.
Jingyuan Liu, Jianlin Su, Xingeheng Yao, Zhejun Jiang, Guokun Lai, Yulun Du, Yidao Qin, Weixin Xu, Enzhe Lu, Junjie Yan, et al. Muon is scalable for Ilm training. arXiv preprint arXiv:2502.1699 16982. 2025.
Yinhan Liu. Roberta: robustly optimized bert pretraining approach. arXiv preprint
1rXiv:1907.11660 1692, 364, 2019.
Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint 11Xiv:1608.03988 2016.
Chao Ma, Wenbo Gong, Meyer Scetbon, and Edward Meeds. Swan: Preprocessing sgd enables adam-level performance on IIm training with significant memory reduction. arXiv preprint arXiv:2412.13148, 2024.
Chris Maddison, Daniel Paulin, Yee Whye Teh, Brendan O' O'Donoguue, and Amaud Doucet. Hamiltonian descent methods. arXiv preprint 1809 050022, 2018
James Martens. New insights and perspectives on the natural gradient method. Journal of Machine Learning Research, 21 2020.
James Martens and Roger Grosse. Optimizing neural networks with kronecker-factored approximate curvature. In International conference on machine learning, pp. 2408-2417. PMLR, 2015.
James Martens, Jimmy Ba, and Matr Johnson. Kronecker-factored curvature approximations for recurrent neural networks. In International Conference on Learning Representations, 2018.
Depen Morwani, Itai Shapira, Nikhil Vyas, Eran Malach, Sham Kakade, and Lucas Janson. A new perspective shampoo's preconditioner. arXiv preprint arXiv:2406.1772 2024.
12

---

Published as a conference paper at ICLR 2026
Son Nguyen, Bo Liu, Lizhang Chen, and Qiang Liu. Improving adaptive moment optimization via preconditioner diagonalization. arXiv preprint arXiv:2502.0748. 2025.
Omead Pooladzandi and Xi-Lin Li. Curvature innooumed sgd via general purpose lie-group precondi- tioners. arXiv preprint arXiv:2402.04553 2024.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, 21(140):-67, 2020.
Thomas Robert, Mher Safaryan, Ionut- Vlad Modoranu, and Dan Alistarh. Ldadam: Adaptive optimization from low-dimensional gradient statistics. arXiv preprint arXiv:2410.16103 2024.
Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning, pp. 4596-4604, PMLR, 2018.
Chongjie Si, Xuchui Wang, Xue Yang, Zhengqin Xu, Qingyun Li, Jifeng Dai, Yu Qiao, Xi- aokang Yang, and Wei Shen. Flora: Low-rank core space for n-dimension. arXiv preprint 1471: 2024.
Yuandong Tian, Lantao Yu, Xinlei Chen, and Surya Ganguli. Understanding self-supervised learning with dual deep networks. arXiv preprint arXiv:2010.00578. 2020.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.1397 2023.
Nikhil Vyas, Depen Morwani, Rosie Zhao, Itai Shapira, David Brandfonbrener, Lucas Janson, and Sham Kakade. Soap: Improving and stabilizing shampoo using adam. arXiv preprint arXiv:2409. 2024.
Alex Wang. Glue: A multi-task benchmark and analysis platform for natural language understanding. arXiv preprint arXiv: httv:8804.00 180407741, 2018.
Minghao Xu, Lichuan Xiang, Xu Cai, and Hongkai Wen. No more adam: Learning rate scaling at initialization is all you need, 2024. URL URLhttps: /attps://aaee /tEtts://axxtv.orrr/bbb///
11768
Zhirong Yang and Jorma Laaksonen. Principal whitened gradient for information geometry. Neural Networks, 21(233):232-244 2008
Jui-Nan Yen, Sai Surya Duvvuri, Inderjit Dhillon, and Cho-Jui Hsieh. Block low-rank preconditioner with shared basis for stochastic optimization. Advances in Neural Information Processing Systems, 36:17408-118
2023.
Yang You, Igor Gitman, and Boris Ginsburg. Large batch training of convolutional networks. arXiv preprint arx:v: www.11788.00000 03888, 2017.
Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh. Large batch optimization for deep learning: Training bert in 76 minutes. arXiv preprint arvi: 1904.00962, 2019.
Haochen Zhang, Junze Yin, Guanchu Wang, Zirui Liu, Lin Yang, Tianyi Zhang, Anshumali Shrivas- tava, and Vladimir Braverman. Breaking the frozen subspace: Importance sampling for low-rank optimization in Ilm pretraining. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025.
Yushun Zhang, Congliang Chen, Ziniu Li, Tian Ding, Chenwei Wu, Yinyu Ye, Zhi-Quan Luo, and Ruoyu Sun. Adam-mini: Use fewer learning rates gain more. arXiv preprint arXiv:2406.16793, 2024.
Jiawei Zhao, Zhenyu Zhang. Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. Galore: Memory-efficient Ilm training by gradient low-rank projection. arXiv preprint arXiv:2403.03507, 2024a.
13

---

Published as a conference paper at ICLR 2026
Pengxiang Zhao, Ping Li, Yingjie Gu, Yi Zheng, Stephan Ludger Kölker, Zhefeng Wang, and Xiaoming Yuan. Adapprox: Adaptive approximation in adam optimization via randomized low- rank matrices. arXiv preprint 11XX1:24003.1998, 2024b.
Rosie Zhao, Depen Morwani, David Brandfonbrener, Nikhil Vyas, and Sham Kakade. Deconstructing what makes a good optimizer for language models. arXiv preprint arXiv:2407.07972, 2024c.
Hanqing Zhu, Zhenyu Zhang, Wenyan Cong, Xi Liu, Sem Park, Vikas Chandra, Bo Long, David Z Pan, Zhangyang Wang. and Jinwon Lee. Apollo: Sgd-like memory, adamw-level performance. arXiv preprint arXiv:2412.05520, 2024.
14

---

Published as a conference paper at ICLR 2026
A RELATED WORK
Optimizer based on structural approximation Due to the desirable properties and convergence of recond-order optimization, various work has been proposed to efficiently approximate Hessian- like matrix, e.g. FIM KFAC (Martens & Grosse 2015) was one of the first work that goes beyond the simple diagonal approximations, and approx the layer-wise FIM Subsequent works extends KFAC beyond MI.P layers (Grosse & Martens 2016, Martens et al.]2018). Further refinement to KFAC are also proposed, including the refinement of eigenvalues (George et al. 2018), fixing the trace (Gao et al. 202T), and refinement by Kronecker product singular value decomposition (Koroko) et al. 2022). Our proposed view point is different from KFAC, which express the FFM usinnn the back- proped gradients and layer input. In addition, KFAC needs to re-derived for different type of layers. On the other hand, our proposed view point is closer to another line of work, aiming approximate the full AdaGrad (Duchi et al. 2011). In particular, Shampoo (Gupta et al. [0188] Anil et al. 2020 is proposed as a Kronecker product approximation AdaGrad. Later, (Morwani et al. (024) explicitly proved that it is a 1-step power interation to optimal Kronecker product approximation. In here, we propose an alternative view as minimizing a upper bound of the approximation error. SOAP (Vyas) et al. 2024) is recently proposed adaptive optimizer that further improves Shampoo based on the connection of Adafactor and Shampoo. In this work, we make explicit connection of those approaches to FIM approximation, and establish the equivalence of structural assumption to optimizers. We also additionally provide connections of gradient operators to FIM approximation, and design new optimizers from this view point. We provide more discussions of our approach to many existing optimizers, including Apollo (Zhu et al. 2024). GaLore (Zhao et al. (2024)). Muon (Jordan et al. 2024), SWAN (Ma et al., (2024). Adapprox (Zhao et al. 2024b). Lars (You et al. (2017). Lamb (You et al. 2019), Fira (Chen et al. (2024a) and AdaDiag (Nguyen et al. 2025), in Sec. [.5 Preconditioning SGD (Li 2017 [Pooladzandi & Li 2024) aims to directly approximate the inverse Hessian through
different structural assumptions.
Comparison to blockwise low-rank approximation (Yen et al., 2023) Recently. Yeen al.
2023) proposed an optimizer based on the blockwise low-rank approximation to the full AdaGrad preconditioner. On the high level, it seems to be similar to our proposed low-rank extension framework. In fact, they are different in several important ways. First, their approach is not scalable due to the use of the full coefficient matrix with shape r X r. where r is the rank. From Section 2.2 in (Yen et al., 2023), the memory cost of their method is (m + n)r². For 130M model with r = 256 and weight size m = 768,7 n = 768, this corresponds to 42 times higher memory than Adam and 84 times higher than the model weight. On the other hand, Alice 77 times smaller than their approach. Our low-rank extension framework also incorporates switching and compensation steps, which are proved be critical for low-rank based method. More importantly, the switching step challenges the common knowledge that one should use the top eigenvectors for a low-rank preconditioner. Our analysis (Theorem (..)) suggests that less important directions should be mixed with the leading ones for better performance.
requires roughly 55% of Adam's 's memory, and
-emory-efficient optimizer Practical training efficiency is a crucial factor when training large models. In particular, there many work focusing the memory efficiency of optimizers, since less memory consumption allows larger batch size, effectively improving the throughput. There are two lines of research: (1) use low-rank approximation to reduce memory of optimizer internal states; (2) remove the internal states, GaLore (Zhao et al.) 2024a), a well-know low-rank optimizer, proposed use Singular value decomposition (SVD) for finding a low-rank projection and apply Adam within it. It can be seen as special case of Alice Fira, an extension to GaLore, add compensation term at each step to turn low-rank update to full-rank, substantially improves the performance. Flora (Si) let al. 2024) uses randomly sampled Gaussian matrix as the subspace to save compute and memory. However, is mainly focused on the fine-tuning tasks. ReLora (Lialin et al. 2023). an extension to LoRA (Hu et al.) (222), periodically merges the LoRA weights to enable full-rank learning. On the other hand, many optimizers requires fewer internal states compared to Adam. Lion (Chen et al.. 2024bb and Signum (Bernstein et al. 2018) introduced optimizers that only requires the storage of I internal states, offering a balance between memory efficiency and performance. Apollo (Zhu et al. 2024), a recently proposed approach, only maintains rank 1 GaLore states for estimating the scaling matrix for the raw gradient. Although it still requires 2 internal states, rank states allows it to achieve SGD-like memory. At the same time, (Ma et al. 2024) developed SWAN, which manages to
15

---

Published as a conference paper at ICLR 2026
completely removes the internal states through two gradient operators: normalization and whitening. and obtain stronger performance than Adam. In this paper, we also show that normalization and whitening operators special cases of FIM approximation.
Comparsion to Sophia Sophia (Liu et al. 2023) is recently proposed second-order optimizer for training LLMs. It requires similar memory consumption as Adam but claims to achieve 2x speed-up. However, (Kaddour et al. 2023) argues that the speed-up is achieved under an unfair comparison. With the fair comparison (i.e. the setup adopted this paper), Sophia consistently performs worse/or similarly compared to Adam, whereas ours (i.e. [Alie RACS achieve near 2x speed-up in both training steps and wall-clock time under the fair comparison setup.
B RACS AND ALICE
ALGORITHMS
Algorithm 3 RACS
1: Input: learning X, scale a, limiter threshold 7, optimization steps T.
2: 8n 0: qo = 0: do- = 0
3: for f =....,, do
4: GV=W.44
5: Obtain St and Qt by Eq.
6:
Bs + (1
Diag St);
7:
q: 3q:-1 (1-3) Diag(( Q.)
8:
G. = Diag, (qr) Diag (8.)=
ア = 7/ max GGlll
9:
oe-1" 73 if > else
Φe = nl|G.|
11: W1+1 = W, - Xna
12: end for
Algorithm
opttimzzer
1: Input: learning rate X, scale a, compensation scale o.,, update interval k. 31, B2, Bg (3g = 0 for
Alice-O). optimization step T', rank loss function L, limiter threshold y, leading basis number
4.
2: Qo Uo 0, Pa 0, d 0, mg O, vo
3: for f = 1..., T do
4:
G. Vw.C
5:
/ = or (L mod K) = 0 then
6:
Q: B,,QQ,1 U. + +1 G.
7:
U, Switch QU, r. U.-1)
8:
else
9:
U. U:-1
end if
11:
U.G.
12:
Q 83Q,-1
+ 3
13:
m, = 31m--
+ +1 -1-6) 1--1)
-2U-11 + +11 3200
15:
w = my
Vol
A Pr, Φ: CompensattionU PU-, 01-1
17: W = W, - >a(Uww + a.D.)
end for
C EXAMPLES OF DIAGONAL OPERATIONS
In this section, we will give some detailed examples on each of the diagonal operations we introduced
in Sec. 2
16

---

Published as a conference paper at ICLR 2026
Example of Diag(.)
This will extract the diagonals of a imput matrix into vector:
Q11
012
d13
M
Q21
d23
Diag( M) = [011, 022, (333
A31
032
433
Example of (-) This simply stack the input matrices sequence into a larger block diagonal
matrix:
M
0
Diags (M1, M2, Ms)
0
M2
0
0
M3
Example of Diago() This will stack the vector element into pure diagonal matrix:
411
Diagy ([a11, 022, a33l")
0
0/2
0:3
Example of D-agx(-)) This will stack the elements in input matrix to form a larger pure diagonal matrix in column-wing wise manner.
0
Diago (I
([au a2
012 a22
1)
021
0
0
0
012
0/22
D BACKGROUND
In this section, we will provide a more comprehensive backgrounds.
D.1 FISHER INFORMATION
Fisher information can also be viewed the "sharpness" the likelihood around the true parameters from the maximum likelihood view point. Formally, under the context of LLMs fe(-), we consider dataset {æ; )- where N total batched sentences, and æ; the token sequence of 7th sentence. One can define sentence-level auto-regressive loss function L = E5=ı c(ri5-1, (æi..)): where is the token index and e is the user-defined loss metric. The corresponding likelihood can be defined as (x..) oc exp(-Xx ( c(Z1,++1, (o))..)))) The standard empirical FIM defined as
Vo log po x V log po(;))
In practice, we often use the mini-batched gradient g Nu Σ log pa (x) with batch size Ng for empirical FIM during training. More discussion can be found in Lin et al. 2024). Throughout this paper, we adopt the notation E' F.
leverage the inverse of FIM to smooth the local information geometry, leading the steepest descent in the probability space with KL divergence metric. This typically leads to faster convergences compared to its parameter-space counterpart; and more stable optimization (Martens) 2020). The update of NGD with step size is
One standard application of FIM is efficient optimization.
NGD
0- XF VoL.
However, square-root inverse is sometimes more favorable than inverse, and has been shown that provides better approximation the geodesic flow, compared with the default natural gradient update Yang & Laaksonen (2008). Empirically, demonstrates stronger performance and desired
properties when using non-constant learning rate (Lin al.) 2024 Loshchilov Hutter 2016
Bergstra & Bengio, 2012 Choi 2019). The corresponding update is to simply replace F with F-2:
0 0 XF VoL.
(16)
17

---

Published as a conference paper at ICLR 2026
In this paper, will use the square-root inverse view point, but analysis is agnostic to it and can be easily extended to the direct inverse. However, matrix multiplication involving FIM and its inverse computationally expensive for large models since F Rmaxm for vectorized parameters E Rem Next, we will briefty introduce Kronecker product and block diagonals, which are two main classes of structural assumptions used in this paper. Their nice properties can significantly reduce the assocaited computation burden.
D.2 KRONECKER PRODUCT AND BLOCK DIAGONALS
Kronecker product, denoted by 69, is to combine two arbitrary matrices into a larger matrix with block-wise patterns. It has emerged as powerful tool for approximating higher-order information like Hessian (Grosse & Martens 2016) or FIM The main advantages are their nice properties regarding matrix operations, leading to efficient practical procedures when dealing with large matrix. We will mainly use two properties:
(A B Ø B
(17)
(A B) Vec(C) Vec BCA*)
(18)
where the first one holds when A,B square-root invertible. The second one is particular useful
to reduce the computation associated with matrix-vector multiplication in Eq. (T). For A @ B E pmaxmmn reduces the computation from Omm²n²) to O(mn² 2 m²n) with _2ax B
Rmxm, and C E Ruxn
For block diagonal matrix, one can also easily compute its square-root inverse:
M. #)
Diaga(M1, M,) DiagB (M,
(19)
where each Mj is square-root invertible. When each M; is also diagonal matrix with positive values, we have the following:
Diage M M.) Vec(C) = Vec
(20)
VMat(v)
where v is the vector containing the diagonals Diage ,M product into element-wise operation.
transforming matrix vector
D.3 OPERATORS FOR GRADIENT: NORMALIZATION AND WHITENING
Recently, there are some optimizers (M et al. 2024 Jordan et al. 2024 You et al. 2017: 2019) that apply operators to pre-process the gradient and use in standard SGD. Empirical evidence has verified their effectiveness in training LLMs. In particular, there are two well-known operators: normalization and whitening, where the above optimizers relies on one or both them. In particular,
G
Norm 1VST
GS
(21)
Whitening -GG")
(22)
where 8 E =; the squared column norm of G. and S is scaling matrix to normalize the columns. Normalizing the rows can also be written in similar format. Whitening operator essentially orthogonalizes G, that compute the closest orthogonal matrix to G under Frobenius norm. In practice, the whitening operator can be iteratively solved using Newton-Schulz algorithm without explicitly computing the square-root inverse.
G2yy with G E mx and S Diag, (s). Namely, vector contains
j==1
D.4 SHAMPOO OPTIMIZER
Shampoo (Gupta et al. 2018) was originally proposed as the second-orde optimization technique over the tensor space. Under the context of transformers, typically matrix parameters are considered. Its
core design principle also aims approximate the FIM with structural assumptions F, R² L m.
18

---

Published as a conference paper at ICLR 2026
with Ru,, R G. and - + G.G. The above update rule be
as the moving average estimate of E|G G.] and G However, the Shampoo paper (Gupta
ctt al. 2018) does not explicitly show why these design choices of R., Lm approximate the [IM In
fact, they only show R @ Lm forms an upper bound of FIM This is not helpful in understanding whether this approximates the FIM or not. Another very recent follow-up work (Morwani et al.
the shampoo preconditioner terms of approximating [IMM] They show the square of Shampoo's preconditioner is equivalent to single step of power iteration of computing optimal Kronecker product approximation to FIM It indirectly establishes the connections FI approximation since the approximation is expressed as an iterative algorithm. To the best of our knowledge, our result is the first to directly establish the connection of Shampoo to FIM approximation minimizing upper bound of the loss with the Frobenius norm (Theorem 3.2
2024) provides an explanation
Nevertheless, the original Shampoo algorithm is summarized in Algorithm 5
Algorithm Shampoo Optimizer
Input: Lm c.m. R., cL,, learning rate X, optimization step T, loss function L.
for t 1,.... do
G, = Ww,C
Lm, Lim,t-1 + G.G.
Rn, RA + G.
W,- Wi-1 AL AL G.
end for
D.5 SOAP/ADADIAG*+
SOAP/AdaDiagt+ Vyas et al. aiming to improve the practical convergence and stability of Shampoo. Their main intuition behind (from the view point of Vyas et al. (2024) is that they show Shampoo is equivalent to performing Adafactor (Shazeer Stern 2018) under the Shampoo's 's eigen-space. Namely, the eigen-matrix of Lm,t and Rn,t t in Algorithm S Since Adafactor is an approximation to Adam, they propose to use Adam instead of Adafactor in Shampoo's 's eigen-space, to further improve the performance. They propose the following update rule:
2024 Nguyen et al. 2025) is a recently proposed adaptive optimizer
m. Bim4-1 + B1) (first moment)
L B3Lm (1 G (Shampoo's Lm,t)
Ru, = Bg + +(1 (1-Ba) Ga) G. (Shampoo's Rn,))
UL, EVD(Lmm.) (Shampoo's left eigen-space)
UR.: = EVD(Rn,,) (Shampoo's 's right eigen-space)
0, 8201-1 + (1-
B,)(UT, (second moment)
U.,.,,UR... U
A = UL
(23)
Vve
These update rules exactly describes the procedure to applying Adam updates in the "rotated" space defined by UL,, and UR.. Due to the computational burden associated with EVD, SOAP proposed to only update UL.. UR, at certain intervals. This leads to the following algorithm:
D.6 ADADIAG AND ONE-SIDE SOAP
AdaDiag++ (gguyen et al. 2025) a concurrent work to SOAP, independently develops the equivalent update rules as SOAP. The only difference is that they disable the [EMA] tracking for Lm, and R The resulting optimizer is both computational and memory expensive due to the storage of Un. Ur and two eigenvalue decompositions. To address this issue, they both propose a one-side version
"Lemma 8 in Gupta al. 2018 Note that our paper assumes Vec(.) is stacking columns of matrix whereas Gupta al. (018) assumes stacking the rows, explaining the reverse order of presentation
19

---

Published as a conference paper at ICLR 2026
Algorithm 6 SOAP optimizer
Input: learning rate A, update interval K, 81, $2, B3, optimization step T.
mo 0: D0g 0
binitialize two moments
Lm.0 0; Rm,u =
>Initialize two EMA states for G
for t = 1,..., .T do
G. = Vw.L
m, 81m4-1 + (1 -
L m = B3Lm + (1- 10 G
DAccumulation for GGT AAccumulation for G G
R = B + (1 Ba) G.
if = or (t mod K) = 0 then
UR. = (
>Get right eigen-space Un DGet left eigen-space UL
UL. EVD (m,))
else
Un URE-1
UL UL
end if
m. UT, -1,,m URR
DGet rotated 1t moment D Compute second moments
0, 12 +(1 -2) n.Un,))
W4+1 = W. - LU me U
voi R,
end for
called AdaDiag and one-side SOAP by only considering either left right eigen-space. The resulting update rule is exactly the same as our proposedOOOOAP (i.e. Eq. (9).
However, they propose this design choice purely based intuition to reduce computation and memory consumption, and do not explicitly reveal the connections to their two-side version. Thus, it lacks the understanding
why two-side version obtains empirically better performance. Based on Sec. 3.4 and Sec. E we show that although one-side version has similar updates as the two- side twins, they different optimizers with distinct underlying structural assumptions. In fact, the structures of SOAP/Adas COAP//daDiag++ strictly generalizes their one-side version, explaning the better empirical performance. The resulting algorithm the following:
Algorithm 7
SOAP optimizer
Input: learning X, update interval K, 31, A2, A3, optimization step T. mo = 0; v0
>Initialize two moments >Iiitialize the EMA state for GG
Qo 0
for t = 1...., ....,Tau do do G, Vw.C
Q: = 33Q,- 1-11 + (1 B
DAccumulate the GGT DAccumulate the first moment
m. 31 f, + Bi
if L = or (1 mod K) = 0 then
U = EVD(Q.)
>Obtain the U
else
U = U
end if
m, UT..m
D-Rotate the first moment Acccuullate the second moments > D-pddate in the original space
v. 8-81-1 (1 (2),
W = W.- - XU f0:
m.
end for
D.7 SWAN
Recently, there is newly proposed adaptive optimizer that completely removes the needs of storing internal states, called SWAN. relies on two processing operators applied to raw current gradient: GradNorm and GradWhitening, replacement of first and second moments. For current gradient
20

---

Published as a conference paper at ICLR 2026
G,
GradNorm(CG G gl
(24)
s1
GradWhiteningg(G- (GG") G / E g)² the standard
(25)
where 9 - Ei= G. is the mean across rows;
deviation across rows, 9 and 8 are m-dimensional column vectors; and 1, is n-dimensional column vector of ones. Then, SWAN performs the following to generate the update:
G GradNorm G))
A GradWhitening
(26)
We can see that the proposed GradNorm is equivalent to normalization up to scaling vn when the mean 9 is 0. SWAN derives these two steps from investigating the L.M dynamics. In practice, SWAN proposes to compute the (GG")- using Newton-Schulz iterations.
D.8 NEWTON SCHULZ ITERATION
In many machine learning applications. like successive whitening and coloring transform (Li et al. 2017; Cho et al. [2019] Choi et al.] (2021), one often encountered the computation of square root inverse of some SPD matrix. One standard approach is compute the EVD and take the square root inverse of the eigenvalue matrix. However, EVD computationally expensive. Another alternative approach is use Newton-Schulz iteration (NS), an iterative updates of two matrix. Specifically,
A
Y.
Zo =
|
Y11 Yi(II
Z.Y.)
Z41
(3I Z.
2
with convergence Y, A3 A and Z. A- v ||| Typically, NS converges very fast with
only steps al. 2018 Huang et al. 2019).
D.9
MUON
Muon (Jordan al., 2024). is recently proposed to speed-up the training of LLMs, that relies on the whitening operator similar to SWAN. The core of the Muon is to orthogonalize the the momentum. The proposed update rule is
m. = 31)
A StradWhitening(/
(27)
Similarly, the GradWhitening step is computed using Newton-Schulz iteration. The main difference between Muon and SWAN is that Muon still requires the storage of first moments as state, whereas SWAN relies on the GradNorm operator applied the raw gradient.
D.10 LARS
SWAN and Muon both involve the whitening operator with hith/without normalization, respectively, On the other hand, Lars YYuu al. 2017) an operator that only relies on layer-wise normalization. For each layer, it simply compute the first moments, followed by a normalization operation. The update rule for each layer
m, = 8,m-- + - G
m.
= ((10))
(28)
|mm
where is scaling function with input of the parameter 0 norn. One major difference of this layer-wise normalization to SWAN is that is applied on the layer-wise level. whereas SWAN applies row/column-wice normalization.
21

---

Published as a conference paper at ICLR 2026
D.11 Low RANK OPTIMIZERS
The primary goal of low-rank optimizer is to reduce the memory consumed by the states of adaptive optimizers. The popularity of low-rank based method for large models starts from the well-known LoRA (Hu et al. (021), where each weight is inserted with an low-rank adapter W + AB with A _ and B R xx during the finetuning stage. This formulation directly modifies the model architecture. Si ct al. 2024) explicitly show that LoRA is secretly a gradient compressor, which translates the modification to model architecture into low-rank optimizer with randomly sampled matrix. At the same time, GaLore (Zhao et al.) (024a) popularizes the use of low-rank optimizers, which demonstrates on-par performance compared to full-rank Adam training. GaLore proposed based the analysis of reversible networks (Tian et al. 2020). However, in practice, transformer may not satisfy the reversibility condition. Thus, GaLore does not provide a clear understanding on why it works LLMs.
Algorithm summarizes the procedures. In practice, GaLore can be viewed as the composition of subspace search algorithm (i.e. SVD) with standard Adam optimizer. The original GaLore does not provide an explanation on the choice Adam. On the other hand, our analysis reveals that the GaLore is approximate low-rank extension different optimizer, OSSAAP AdaDiag/one-side SOAP, that is generalizes Adam (see Sec. 3.4).
Algorithm 8 Galere Optimizer
Input: learning rate X, decay rates 31, 82, rank r, update interval k, scale a fort t 1...., ,T do
G, = Vw,
if / or mod k = 0 then
U, = SVD(G,,T)
else
U, = U,-1
end if
G,
A Adam ot, (1, ,1, 32)
W. = W. + A
end for
Since the states of Adam optimizer are based on the projected gradient ot, the overall memory consumption of GaLore is mn 2nr mr.
D.12 APOLLO
Concurrent to this work, there is a recently proposed optimizer, called Apollo (Zhu et al. 2024), that only scales the raw gradient and obtains SGD-like memory with Adam-level performance. They propose to scale the columns or rows similar to normalization in SWAN, but the scaling factor is estimated following the procedure proposed by Fira (Chen al. 2024a). The core idea of Apollo is to obtain A from GaLore algorithm (Algorithm 8]. followed by computing the column/row- norm of A. This norm will be used as the scaling factor for the raw gradient G. Apollo has many variants. In particular, we choose Apollo-mini and Apollo-svd, where the former uses rank-1 random projection for scaling estimation, and the latter relies on the use top r singular vectors as the projection, same GaLore. Apollo-mini only maintains the rank-1 states, leading to significant memory savings. The memory consumption is mn + 2n + 2 for parameter W E m= And Apollo-syd consumes the same memory as GaLore.
the procedures.
Note that when rank is to 1, they propose to use global scaling scaling-
| instead of row/column-wise |
D.13 SUBSPACE ITERATION
The subspace iteration method-ake known the block power method is classical iterative technique for computing the dominant eigenvalues and corresponding eigenvectors a matrix. generalizes the basic power method from operating single vector operating subspace,
22

---

Published as a conference paper at ICLR 2026
Algorithm 9 Apollo Optimizer
Input: learning rate X, decay rates 31, 82, rank r, update interval k, scale for f = -I...., .....Td do
G. = WW.
if = or mod k == 0 then
U. ~ (0,
seed +- independent new random seed
else
U U1-1
end if
a: U.
A
Adam( (1,82)
S. Diag. (rag((sn: (80..... 5m) {st
[Ar.alla To
W. Wi-1 + XaG,S,
end for
typically spanned by a initial matrix. When the initial matrix is closed to the targeting eigen-matrix, the convergence is fast. Empirically. we found that only step of iteration enough to give satisfactory performance. Algorithm (10 summarizzs the subspace iteration algorithm to for finding the top r eigenvectors of matrix A. We can see that the computation is bottlenecked by the matrix multiplication AU; which is ((m²r) if only performing step.
Algorithm 10 Subspace iteration
Input: symmetric matrix A E mRm iteration step T, initial matrix M E m
Uo = M
for t ,,...,T' do
H. -AU,-1
U, QR decomposition H,)
end for
V UF AUT
U = EVD(V)
E SOAP: A FURTHER GENERALIZATION TO OSOAP
All previous structures, apart from the one behind Shampoo, are under the class of block diagonal. However, this only allows one to model the off-diagonal elements near the diagonal. Structure under Kronecker product, like the one behind Shampoo, can go beyond this. Therefore, we can consider combining the structure of OSOAP with Shampoo, to obtain the most general structural assumption in this paper. We show this exactly recovers SOAP (Vyas et al. 2024).
Specifically, we consider the following structural assumption: H {(Un @ UL) 2 UL)"), where UR E Raxn U, E pmxm are orthonormal matrix, and D E R pang@m is a diagonal matrix with positive values. We can easily show that structure behind OSOAP special case by constraining UR = In; and Shampoo also a special case by constraining D to be decomposed by Kronecker product (refer to Sec. I.T.
Similar to OSOAP it is hard to directly minimizing Eq. (2) with this assumption. We can approximate the solution by 1-step refinement procedure OSOAP
Theorem E.1 (I-step refinement of SOAP). Assuming the above structural assumptions. Consider the following I-step refinement: (1) assume D can be decomposed as Kronecker product of two diagonal matrix, find corresponding Un. UL: (2) fix Un. UL. fid corresponding D. Step (1) admits
analytic for minimizing the upper bound of Eq. (2) (i.e. Eq. (2)):
Un EVD G)),
U. = EVV
23

---

Published as a conference paper at ICLR 2026
Step (2) admits analytic solution for minimizing Eq.
D = (B(U)
The proof is straightforward combination of Theorem 3.2 and Theorem 3.4 and can be found in Sec. G.7 One can show that the corresponding optimizer associated with the above result is equivalent to SOAP (refer to Sec. F.4 for details).
DERIVATION OF UPDATE FORMULA
In this section, we will explicitly show how to connect the solution from minimizing reconstruction loss of FI (Eq. corresponding update rule.
F.1 SHAMPOO'S 'S UPDATE FORMULA
The key update formula of Shampoo
W = Wi- + AL A.
Proof. From Theorem 3.2 we simply apply the properties of Kronecker product to square-root version of natural gradient descent:
Mat (
Mat
(RA
G.R ))
MMat Vec
-m G G.R.
F.2 GENERALIZATION WHITENING AND NORMALIZATION
The square-root NGD update with F in Eq.
(3) in Theorem 3.3
Mat
F (
GGE
(29)
Proof. From the solution in Eq. (3. we can simply apply the properties of Kronecker product as in the derivation of Shampoo's 's update:
Mat ( (g)
Mat.
(In 8 M)-
Mat Veec (VTM &G)
= |
Similarly, the square-root NGD update with F - S ( Im is
Mat r ( g
3
(30)
24

---

Published as a conference paper at ICLR 2026
Proof. This is trivial by applying the property of Kronecker product:
Mat((S @ Imm =g)
= -Voo(V
-GS-
F.3 UPDATE FORMULA FOR FOR[OSOAL OSOAP
Proof. From the Theorem
3.4
we can apply the properties block diagonal and Kronecker product
with full-rank U:
Mat
F
g)
M 12
...., ,,.,,M," : ) g))
Mat
Diagg
Mat
Diagg
(UD UT ..... UDi
-0=000 UD, g)
Mat.
(In U) Diago (v D Dn) ( [g)
MMt
(In U) DiagB( D
D Dn) Vec GU*
UTG
-Mat.
(In U) Vec VE[(UT]]]]
UTG VE[(UTG)2]
= =Mat
Vec
UTG
=U
F.4 UPDATE FORMULA FOR SOAP
Based on the Theorem E.1 derive the update formula of the corresponding square-root NGD following the same procedure as OSOAP
GUR UU.
Mat
f (
-U.-
|E(UUEGU GUR)
Proof.
Mat
F (-
MMat
(Un @ UL) Diagm (E (
-Mat
(UR 8 Ur.) Diagm Diag (E(U, Vec (U GUR)
GEGUR
Mat (Un @ Un)
Vec
VEIUUGGURY
GTGUR Ut
MMa
Vec
UL
VE(UTGURR))
GTGUR Uh
-UL
V
25

---

Published as a conference paper at ICLR 2026
Therefore, design the optimizer based this update formula and exactly recovers the SOAP's procedure (Eq. (23) in Sec. D.5
G THEORY AND PROOF
To prove the results, we need to first introduce some useful lemmas and inequalities.
Lemma G.1. Assume F is block diagonal matrix with squared block matrix M; E Rmxx,, then
min F
F
min (M.)
Σ 2TM +C
(31)
where g: is the thh column of gradient G, C is constant that is idenpendent of F, and F is the FIM
Proof. This is straightforward by expanding the Frobenius norm (F-norm)
| - Fl?
Tr (F FF)" F))
=1 - 2Tr (f F) +C
F F.. F.,F++CC
Σ M 2Th
+C
where Fl. Rma indicates the tth row vector of F and F. is the 1th column vector: The last equation obtained by the fact that F is a block diagonal matrix. So only the values of at the position of non-zero values F contributes to the trace, which is exactly the outer product: E E'g.g?"..
Lemma G.2 (Powers-Stormer inequality). For positive semi-definite operator A, B, we have the
following inequality
Tr((A- - B) B)) I A -B²111
(32)
where || [|a is the trace norm.
G.1 PROOF OF THEOREM
Proof. From Theorem G.
we have
F F
M 13 2 (M,
L8 [gggg.") 1
m
Mass 2(92,)
1=1
By taking the derivative w.r.t Mijj. we have
Mijs = -k)s=E1o7,1
Thus, have F Diag( (E(gg)))))))
26

---

Published as a conference paper at ICLR 2026
G.2 PROOF OF THEOREM 3.2
To prove this theorem, we need to leverage the Theorem 3.3 for generalized whitening (Eq. (3) in
Sec. 3.3 This is proved in Sec. G.3 But the following, will provide alternative proof for completeness.
Proof. From Theorem
G.1
we have
F
Fl
M 1 - 2 TrrMMMMMM.,,, +C
-n M || 2TT(MM
9:g 1)
=n M p 22TMM" +
To minimize this, we take the derivative w.r.t. M, we have
2nM 2E[GG" M EGG
Next, prove another proposition that is "symmetric" to the whitening results in Theorem 3.3
Proposition G.3. Assume H. = RR, & I where R., E Rax* is SPD matrix, then Eq. can be analytically solved with the optimal solution as
1 G
R
(33)
Proof. Since R., & Im does not have a nice block diagonal structure like the previous proposition, need to analyze it a bit more. First, have
Rn X Im - FF??
Rm & 1m - 2Tr
RRn @ Im)² Elgg")
+
Since we only care about the diagonal of Z, therefore, we only inspect the block diagonal of Z with each block Z, of size I mxx. and 1...., n. By basic algebra, we have
Z:
Ragag
where gt is the t column of G. Therefore, we can simplify the trace of as
Tr(Z) Tr(Z,)
i=1
Tr
R R
Ra G [
27

---

Published as a conference paper at ICLR 2026
where JG is the element of G at 7th row and 8thh column. Now, let's perform the same analysis of the following quantity
(Im @ R E[G
where [ the vectorized transposed gradient G Namely, now stacks the rows of G instead of columns of G like g. This object is simple to treat due to its block diagonal structure, by algebric manipulation, we have
Tr (I 1 Rn,
R.
Tr(Raj
G²
kihh colurn G
over blocks
22
Σ
Rug G GGg)lalCe (G)
Now, let's change the variable = i = hi and h: = j. the above becomes
(Im @ Rn MIGGGG
TF
m
n 10
Rak GT
j=1
= -Z)
(34)
We should also note that
RR, @ Im
Tr ((RR, & Im) (Rn @ Im))
.Tr ((R,, R.)) Im)
=Tr(R.), R. Tr Im)
TTr ((Im X R.,
(Im & Rn))
=II(Am & R )
Therefore, by using the above equation and Eq. (4), the original minimization problem is translated
arg min Rr, & Im arg min Im @ R. - E - G
FI2
R
R.
Thus, can leverage Theorem 3.3 obtain the optimal solution
R" = 1 L G]
With the above two propositions, start to prove Theorem 3.2
Proof. First, note that
| X L FI
=|| RR, & Im))
(I, & Lm)
Elgg"
B
A| - -CC
28

---

Published as a conference paper at ICLR 2026
Next, we will upper bound this quantity. First, we have
AB - CC AB-CC-ABBCC)(A-CCT
By triangular inequality, we have
AB CC
< - -
< |
S (I -
|| C| ||||+
C -
Now, the squared norm can be upper bounded by
|B - CC <3 (I4 - - + -
C || B CI
<3 A mn
Vmn C 14 -
(35)
(36)
The first inequality is obtained by the fact that for any three matrix P, Q and H, we have
|P + +++++HI
- (1) P || + H r)
=| Q1? H 1 1 + 2|||P||a|| 2
33 (1) | -H
The second inequality is obtained by directly applying Powers-Stormer" inequality and Holder's s inequality. For completeness, we will show how to upper-bound || A C the rest can be bounded in the same way. From Theorem G.2 and both A, C' are SPD matrix, we have
| C - ||A² 1²
Then, we can select p = = 2 for Holder's inequaity and obtain
|| C²1 /mm|| A C²Il
where /mn comes from the || in Holder's inequality. By substitute it back, we obtain the upper bound.
We see that minimizing the upper bound Eq. (36) is equivalent minimize each | C² B CfIl individually, and
1 RR., @ I. F|F
_
C2 I @ Lm FIlF
Thus, from Theorem 3.3 and Theorem G.3 we prove the theorem.
G.3 PROOF OF THROREM 3.3
AND THEOREM
Instead of proving the Theorem 3.3 we propose a generalization to those gradient operations, where
Thcorem 3.3 is special case.
Structure assumption We consider H = {S 8 M) with identical SPD M E m and positive diagonal S R The following theorem proves that the optimal solution can be solved by fixed-point iteration.
Theorem G.4, Assuming H. = {S & M) with positive diagonal S R and SPD M e mm and EGG (G" G')²] contains positive values, solving Eq. (2) admits a fixed point procedure:
Diag(S) ( MG))
M E
(37)
M M
|
The solution Diag(S*) converges to the principal eigenvector of ( G') up to a scaling with unique S* 6 M
29

---

Published as a conference paper at ICLR 2026
To prove Theorem G.4
first introduce some classic results.
Theorem G.5 (Perron-Firobniius theorem). For matrix A E guxa with positive entries, the eigenvalue. The corresponding eigenvector of A is called Perron vector and only contains positive components: Av r with U, > 0. In addition, there are no other positive eigenvectors of A.
principal eigenvalue is positive, called
Definition G.6 (Hilbert projective metric). For any given vectors v, in C/(0} where C' is closed convex pointed non-ne non-negative cone C, i.e. Cn(-C) = {0}. the Hilbert projective metric is defined as
U/ Wig
min R W,
du (v, w) log
max
log
This is pseudo metric since it has scaling invariance property: du (n, am) = du (v, m) for O. This means dæ(vv, m) = 0 does not mean v m but = am with some positive scaling a. However, this is metric on the space of rays inside the cone.
Theorem G.7 (Birkhoff-Hogg theorem). Let P E Ruxa be positive matrix and let
Py) S adæ (æ, y), Yx, y E C+, x y) where C+. is the cone that each element is non-negen non-negative and is the induced equivalence relation. Namely, if x y, there exists a, B such that oar y < Bar, and x y means y C++. Then, it holds
<(P) inf {a 0:
1 (PP) with A(P) = max R.R.R. PaPaj
PiPPA
R PP) = tanh 4
This theorem suggests that when P is a positive matrix, the corresponding linear mapping is contrac- tive since tanh(-) S under Hilbert projective metric.
Now, let's prove the Theorem G.4
Proof. First, we can simplify the Eq. using Theorem G.I
S M FI
(T. + C
Then, simply take its derivative w.r.t. s;, and obtain
25, M 2 (M
S, Tr(ME[gg.
| 1
Diag(S) = Diag (E MG])
M
Similarly, we have
Σn 1=1
S.E.g.
M
S
|
These define iterative procedure. Next, will show it converges. Let's substitute M into S, and obtain
S = Diag
(Eg Eo (G''G a(S)
= =Diag
EGg
GG SG G
H
30

---

Published as a conference paper at ICLR 2026
where a(S) is the scaling term. In the following, use E as Boar. Since show
a
S E
Σ S
we can write S in its vector format:
= E
H 8.
From assumption, we know P contains only positive values, let's define a quotient space for positive vectors and q under the equivalence relation s' if as' for some positive scaling Namely, we define a space of rays inside the positive cone. Therefore, the Hilbert projective metric becomes real metric inside the quotient space.
From the Theorem G.7 we know the linear mapping associated with P is contractive. Therefore, we can follow the proof of Banach fixed point theorem the previously defined quotient space with Hilbert projective metric to show the convergence of this fixed point iteration on 8.
Now, show the solution 8* is always positive. Since it is converging, therefor, the solution satisfies
a(s") Ps*
This equivalent to finding the cigenvectors P. By leveraging Perron-Frobenits theorem (Theorem G.5). we know s* is the principal eigenvector of P. and only contain positive values. It also easy to verify that this fixed point converges upto a positive scaling factor (this is expected since the contractive mapping holds true for the quotient space with Hilbert metric, that is invariant to scaling.)
Although s* is unique, but S & M is, since for arbitrary positive scaling 3
1 L
38 = M
S M S" & M
Therefore, Theorem 3.3 is direct consequence by substituting F I, @ M and F = =S@ I. into
Eq. (15).
Next, we prove Theorem 4.1
Proof. From the Theorem G.4 the iterative procedure for Q can be simply obtained by taking the
diagonals of M:
Diag (E [GSG]]
||
Following the same proof strategy of Theorem G.4
we substitute Q into the update of S and re-write it into the vector format. First, let's rewrite the update of S
S E
1=1
PT
||g|
where P E[G Similarly,
Ps Thus, [ell
a=a(s) P'p'P'P Ps
From the assumption P contains only positive values, we can follow the exact same argument made in Theorem G.4 show the convergence of this fixed point update and the positivity of the final solution s* Precisely, s* and q" are the right and left principal singular vectors of P. respectively, and S Q" unique.
31

---

Published as a conference paper at ICLR 2026
G.4 PROOFS OF[OSOAP
G.4.1 PROOF OF THEOREM 3.4
Proof. For simplicity, we omit the subscript f in Uy. If we assume all D; are equal and only contain positive values, then each block U are the same for all i, and it is SPD matrix. Then, to minimize the the loss Eq. (2). we can directly leverage the whitening results in Theorem [333] and obtain M = E Due to the structure of U UDU", the optimal U* is exactly the eigen-matrix of M
Next, we prove for any fixed U, we can find the corresponding optimal D, From the block diagonal structure and Theorem G.I
we have
F F
UD 2Tr
+C (U"E[g.'(]
Dill - 2Tr
Ty"Elgg UD
+C
Hi
D
is u H
1=1 j=1
Taking the derivative w.r.t. Di.jj. we can find the optimal Di.s is
Dijs =u- H.u
[u 9
Now, by simple algebra manipulation, we have
D = M (E g:)))
where Diagm is expand the vector to diagonal matrix. Finally, for D, we have
D Diag (B((UT G)
(38)
The optimality of D can also be obtained by leveraging the Lemma in (George al. 2018). and set the eigenbasis as I, @ U.
G.5 PROOF OF THEOREM G.8
Here, we present the following proposition that analyzes the issue casued by using low-rank tracking Q compared to the true Q:
Proposition G.8 (Subspace switching). Assuming the setup mentioned above and all assumptions OSOAP are satisfied. We further assume the low-rank U mamxrr is obtained the beginning i I time block by EVD(Qi where Qúx is the true tracking state. Further, we assume the stability of the eigen-basis such that gradient Ge during + time block shares the same aagen-basiss as Qix Then, we have the following for true statistics Q14417 at the end of the + time block:
(i+1)b
(x+1)k
Q(11IAA
Σ G.G.
Σ G.G. +UE,UU? t=i++1
(39)
1=i++1
where G, = Uo is the low rank reconstructed gradients, Σg R (mm-))x is diagonal matrix with positive values, and U, the complement eigen- such that [U, U.] will form the complete eigen-busis of Q
Proof. Within the time block i + with low-rank mapping U, the gradient at each step can be decomposed as
G, = UU" + (G. UU G.)
32

---

Published as a conference paper at ICLR 2026
Therefore, the true state ((11) be simplified
(++1)k
Q(+1)* =
Σ G.G.
t=i++1
(++1)k
Σ
(G, +
R₂)
tik+1
(6+1)
Σ G.G.
+ G.R. + R.G.T +R,
tik+1
The third equality obtained because we assume G. G.
shares the same eigen-basis basis as Q Namely,
A 0
U
G.G. -UU,U 0
Σ,
U
-UAUU" +UEE,
where A, and Eg are diagonal matrix. Then, we have
G.R, -UU.G.
-UU" +UEUUU -UA,U"U,US U
In addition, we can also simplify
R
-U..G.G
(UA," UUE.U!
-UEU??
Therefore,
(i+1)k
Q((t))N = Σ द. U.E.U!
t=th+1
G.6 PROOF OF THEOREM G.9
Here, we prove the following theorem to analytically search for the optimal scaling S, by minimizing the reconstruction error under Frobenius norm:
Theorem G.9 (Optimal compensation). Assume that the conditions of OSOAP are satisfied. With the proposed form of compensation (Eq. (4 minimizing FIM reconstruction loss
||S, @ ) Fall
admits analytic solution:
Vm-
Diag( (S.)
(40)
E IT ITTTT
where 1m m 1, E Th" the column vectors with element and is element-wise square.
33

---

Published as a conference paper at ICLR 2026
Proof. For simplicity, we ignore the subseript t for the following proof. First, we let O Shen, the loss function can be written as
10 2 U.U - F.li.
1O (.UUUTIB--TTQOUUUUYYUUUUU
Σ
M- m
ΣΣΣ
U2 M..
where M, = Diag(E[((U ()), 9i is the ht column of G, and Uoek is the element in j row, hth column of Ue. Then, take the derivative w.r.t. Oze and we have
20
U2 Jjk =
Σ
B gb)
E
(UT9))
m
This form still requires the access U.. Next, let's simplify it. First, let Ú [U,U.] to be the complete basis, we can show
Σ g.) (T)
k=1
-Tr((UT g g T gi)
-Tr(g ūūT gi)
- g:
Now, let's re-write the above in different format:
Σ (Ū g.)
k=1
TÜO g.)
-Tr(O" 9:9, ŪÚ)
-Tr ([
U"
U 9.9 u.g.
(I(UU"((gg))) + U..'"
-r((UT g.) TU) 9.)) + 9)
g)) Σ U.
k=
k=1
Therefore, we have
E
U 9
Eig. 9,
U g.)
k=1
So, have
E G2
(U
Diag(O) (O)
mm
34

---

Published as a conference paper at ICLR 2026
and
m G (UTG)
Diag( D)
G.7 PROOF OF THEOREM E..
Proof. The proof strategy is straightforward combination of Theorem 3.2 and Theorem 3.4 First, when we assume D has the Kronecker product structure, one can easily write
(UR (L) (L) & UL)
=UU UL) (S @ (SLU))
R x (UIS
B
Therefore the loss (Eq. becomes
AB - CC
where [lgg" 2. This exactly the formulation used in Theorem 3.2 with R URS and L U.S.U.
Thus, by directly utilizing Theorem 3.2
we can see the optimal solution
UnS =
U = E[G
Due to the structural assumption of Un. Ut. Sa. St. their corresponding optimal solution can directly obtained using eigenvalue decomposition.
Now, let's prove the optimal D with any fixed Un U, This is also straightforward by applying the
same technique as Theorem 3.4 The loss can written as
(Un 2 UL) D Unn s Un) FI?
III
=|[LiDIT 13 21r 1 gg -
+C
Since it is easy to verify orthonormality of IT, iT I, the above is simplified to
D | 2Tr II' [[gg""][111
= Σ D3 - D Du Inn Elgg" [1]
i=1
where [II] is the 7th column matrix 11. Then, by taking the derivative, the optimal Dac
D [m Elg Vec"
=E[(
Therefore,
Diag D*) = E[TT
=E[((UT Up)
=Vec((((U
+ D Diag (BB(U))
35

---

Published as a conference paper at ICLR 2026
H CONVERGENCE OF ALICE
In this section, we will prove the convergence offAlice under the continuous-time setup. We initially want to adopt the convergence proof techniques from GaLore (Zhao et al., 2024a). However, the major drawback of their approach is the assumption of the fixed projection U. In practice, one will not fix the projection, and instead, it is crucial to switch different U to achieve satisfactory performance. Therefore, this mismatch requires to find alternative technique that can handle the Lime-varying U. (Maddison et al.) 2018) proposed an elegant framework for analyzing the behavior of optimizers by converting the optimizer updates into continuous-fime ODEs and examining suitable Hamiltonian function. This approach has been recently adopted to analyze the optimizers' convergence behavior (gguyen al. 2025 Chen et 2023
Let's first write down the discretized update of Alice
W. WiI - (U Vite
Mr
G. - G
Vm Vm
2 +E
M, M1 + 101-8110 -
V. = B2V2-1 + +(1 -1- 11=:)(U'G)))
P: = B1 + (1)mm G GU
H + (1- mm)mm) 1 (G. UU"
where is the effective learning rate, U E Rm is the low-rank projection, A A and v are elementwise operations for any matrix A, B.
Next, we can convert the above discretized update into continuous-time ODEs:
dW. U.
M.
Vm Vm G. U,U, G.
dt
VV, -E
D +
dM. U.G. M
dt
dV.
-V,
dt
d
dt 1m 1m (G. - G.)? c
U. = Gt)
(41)
where T(. :) can represent mapping from gradient G, (and the history of gradients) to the low-rank projection matrix U,.
To show the convergence, need to design key quantity called Hamiltonian HWW, M., Vi) such that it satisfies:
min H(W, Mc. Vi) = min L(W)) W.W.V.
(42)
W.
where L(W,) is the loss function we want to minimize. This property ensures that the minimization of is equivalent to the minimization of the actual loss function.
With this property, if dH dr < O, can show the Hamiltonian is non-increasing, and it will reach to
local optimum (i.e. dH
0) at the end. We analyze the state W., M, and V. under dH
and if this leads to G; = O, we can prove the weight ODE trajectory will reach to a local optimum of the loss function L(W,).
We propose the following Hamiltonian:
H(W, M. V.) L(W)) 1
M VV.+
M,
(43)
where (A, B) = Tr(A" B) represents the matrix inner product. Before the proof of the convergence, let's first prove lemma.
36

---

Published as a conference paper at ICLR 2026
Lemma H.1. Assume matrix A has constant columns with positive values, then for any positive semi-definite projector P, we have
PG
G 2 0.
(44)
A
Proof.
PG
G
A
PG
ik G
A
(PG) ik
Gik
G
PG 2 0
ak
where ag is the value of the k column of A and G2 is the tthh column of matrix G. The last inequality is from the definition positive semi-definite projector.
Next, we prove the convergence of Alice
Theorem H.2 (Convergence of Allic).. Assuming the ODEs defined in Eg. 41) satisfies assump- lions of Alice with the Hamiltonian defined in Eq. (43). it satisfies
dH dt
20
(45)
and dH 0 implies G;
aL(W.) BW.
0.
Proof. First, we write down the form dH. dt
dH
OH
dW: OH
dM,
OH dV. OV. dt
dt
OW,
dt
aM,' dt
G. U. VV, +E
M
Vm
G: U,U,G. G
VV.+c U
M.
G. G. Mi
V +c
M
VV(VVV++ (U?G..2 -V.
M
G: -U-UGUG
M. UTG.
G. VV.++
Go, Vm
Vortt
VV.+
M
//
1
M
G1)2 -v.)
M.
V
4
V.lVV. +E72
G. -.--U U.UT
Mi M, VV+E"
M2
Vm
G.
G.)²
D +
VV.(VV + e
Term
Term
1
M
V.
4
37

---

Published as a conference paper at ICLR 2026
It is clear that Term 2 is non-negative. So, first examine the remaining terms (i.e. other terms
except Term I and 2). Let denote
Mr viite
θ for simplicity:
1
O M.)
Vi
4
VA
Mi
M Viik
M
Visn +e)²
V Vi,ik te
m
eM
V - e)² 0
Now, let's examine Term 1:
G. - -.U,, Gi Gi, + - Gi: (c(
U. G.
v + 2 0
where Ue [mmx(m-r) is the complement basis of U such that [U, U²] will form the complete orthonormal basis. The last inequality from Theorem H.l with p= U,U?, A v 4, + e. also trivial to see 4 contains constant columns with non-negative values.
Thus, combining the above, we prove di S 0. Next, we consider the behavior when Hamiltonian
reaches the local optimum. When dh
d 0, we can immediately see M. 0. U G. 0 and
(Ge, G.-U,U,GG V 0.
We examine the last term:
G. - G.
Gt,
V
GUIUU.U ik
=0
V tc
GUlUU. G)ik
Ai +e
V +
G
G G G)
$ Pi +
ik
=0 0
Lk
+
G
Therefore, the weight trajectory W. will reach to local optimum of the L(W.) when dH
I FURTHER DISCUSSION
I.1 CONNECTIONS BETWEEN DIFFERENT STRUCTURAL ASSUMPTIONS
Here, we will make explicit connections between different structures in terms of their generality.
OSOAP
generalizes Adam Since the structure behind OSOAP is Diags U UDDUU ..... UD,U when constraining U Im. the resulting structural is pure diagonal matrix Diagn(D,..... (D....,, D.). This coincide with the pure diagonal structure behind Adam.
38

---

=(Un & UL) DD & DL (Uk o U.)) Published as a conference paper at ICLR 2026
OSOAP generalizes S & M
OSOAP not only extends Adam, but also generalizes the structure considered Theorem G.4 Consider setting D; = S,D, then we have
U = S; D
M
Therefore, (UD, normalization and whitening special cases of Theorem G.4 OSOAP generalizes these two gradient operators as a consequence.
Diage [S M M, Since the structures of
SOAP generalizes OSOAP
We can re-write the structural assumption OSOAP the Kronecker
product format:
Diage (UD,UT ...... ....,UD,U')) (I U) Diags DD...... D (I
D
is clear that this a special case of SOAP structure by containing Un = I.
SOAP generalizes Shampoo The structure of Shampoo consists of SPD matrices R. and L.-- If we eigen-decompos those, and let R. = URDnUU and Lm = ; then the structure of Shampoo can be re-write as
R X LM
=UURV (ERVDaUTD (UL/
This coincides with SOAP's structure when the positive diagonal D can be decomposed based Kronecker product.
1.2
MEMORY CONSUMPTION COMPARISON
Apart from the Table I provided in the main paper, we also include the memory consumption of more optimizers.
Adam
Gal.ore
Fira
Alice
Alice-O
Total
3m
mi + 2nr mr
2nr + mr
2nr mAr
2ur + m71
Weight
123
771
17222
1723
17233
First moment
170
73
13T
Second moment
TU"
13
13T
N/A
1713"
1723"
C.
N/A
N/A
N/A
Q
N/A
N/A
N/A
N/A
Table The memory consumption of low- tow-kank optimizers. Here, we assume the weight has shape m x n with m < n and < m. Note that memory consumption and r is typically very small. output layer. For LLaMA, we select r 512, leading to y2 x 262 K, which is 8x larger than n. However, both are marginal compared to mr = 5120 512, which 10x larger than 2 and 80x larger than n.
For example, let's take the largest possible n 30K
1.3 COMPARISON TO PREVIOUS FIM APPROXIMATION
The idea of efficiently approximating the FIM is not novel, and has been extensively studied in the previous work. KFAC (Martens & Grosse 2015) is a well-known second order optimization algorithm for neural networks based on structural approximation to FIM In particular, they explicitly express the FIM in terms of the gradient w.r.t this layer's output and the input to this layer. Namely, they directly decompose the gradient G used in our paper, whereas in our paper, treat G whole quantity. In addition, the original KFAC also makes one crucial approximation: E & B] ss E[A] & B They do not show theoretically whether this is good approximation and under what metric, but
39

---

Published as a conference paper at ICLR 2026
argue in practice this gives good accuracy. Last but not least, the original KFAC considers the FIM for entire layers, and each block (i.e. corresponding to each layer) under their setup is actually the entire FIM aim to approximate. To be precise, the principle is to structurally approximate the diagonal block of KFAC without decomposing the gradient G using KFAC.
Another more related work is AdaGrad (Duchi al. 2011]. If we only considers layer-wise gradient, the full AdaGrad matrix the Σ gigh Although our principle to approximate FIM in practice, we use EMM as the approximation of I.. Therefore, our principle can also be vIwed as structural approximation to full EMA AdaGrad matrix. Under this view point, there are some previous work on structural approximations. Shampoo (Gupta et al. 2018), was originally proposed as an approximation to full AdaGrad matrix under the online learning setup. However, they do not explicit show under what metric this approximation is based and whether it is optimal or not. Later, there follow-up work (Morwani et al. 2024), showing that Shampoo is 1-step power iteration algorithm of optimal Kronecker product approximation to FIM under Frobenius norm (Koroko et al. 2022). In this paper, we further extend the idea of approximating AdaGrad matrix more efficient structures, under Frobenius norm.
1.4 SOLUTIONS GENERAL BLOCK DIAGONAL STRUCTURES
Here, we present the solution of the most general block diagonal approximation to FIM and discuss why this is not practical.
Proposition L.1 (General block diagonal approximation). Assume = Diag( M,,..., Mn) with
SPD matrix M; E Rmxm, then minimizing Eq. (2) admits analytic solutions:
M = Elg.9?1
(46)
where gi is the thh column of G.
Proof.
This is straightforward by taking the derivative w.r.t M,. By leveraging Theorem G. we
have
| F12
M. - 2TMM
Then, taking derivative w.r.t. M,, we get
M = =lg.9"
From the above, although it admits analytic solutions, its corresponding practical procedure requires the EMMA to estimate [[gg E Rmxm for all i = -,..., n, leading to expensive memory cost nm? Another issue that does not allow efficient computation of the inverse. From the property of block diagonal matrix, to compute --1 g. one needs to invert each M,, incurring ((m²) computational cost. In total, the computational cost of inverting matrix F is O(nm²). Due to both memory and computational constraints, this structural assumption does not lead to practical algorithms.
1.5 CONNECTIONS EXISTING OPTIMIZERS
Lars and Lamb Lars and Lamb You et al. 2017 2019) relies on the normalization operation of the gradients. The main difference that Lars proposed to normalize the raw gradient, whereas Lamb normalizes the processed gradient from Adam optimizer. They also involves a scaling parameter that a function of the weight. Compared to the normalization discussed in Sec. 3.3 the main difference is that we use channel-wist normalization with unit column or row norm, whereas Lars/Lamb uses matrix-wist normalization where the norm of the matrix is regularized to be 1. However, if one vectorizes the matrix weight into vector, and stacks those vectors into larger matrix. Then, Lamb and Lars can viewed as a 1-sample approximation to FIM under the structure considered in Sec. 3.3
40

---

Published as a conference paper at ICLR 2026
Muon Muon (Jordan et al.) 2024) performs the whitening operation the momentum. Sec. D.9 gives a brief introduction. In fact, Muon can be viewed as a special case of Eq. (3) in Theorem 3.3 by
considering the following approximation:
E B|GG") = E|G]
(47)
Thus, the resulting operation becomes the whitening the [[G] which estimated by the momentum
in practice. One potential mismatch of the gradient whitening from the original Muon (Jordan et al. 2024) compared to Theorem 3.3 is the lack of the coefficient vn. preventing the parameter update to adapt to the shape of the weight matrix. Later, (Liu et al., 2025) found that this term is critical for large-scale LLM training, and add it back to further improve Muon. Their argument is from the empirical intuition match the update RMS norm of Adam, whereas Theorem 3.3 proved that this is indispensable from a theoretical perspective.
Another advantage of Muon is its relatively low memory consumption compared to Adam since it only requires the storage of first moments. However, our proposed Alice uses less memory than Muon. Let's consider weight matrix with shape m m, and Alice uses rank 0.25m, which is consistent to the setup for experiment. Thus, Muon will use m² memory for state and Alice will only consume m² 20% less memory than Muon. 16
SWAN SWAN composes the gradient normalization and whitening operations as the replacement of Adam's first and second moments. Each individual operations can viewed as a special case of Theorem 3.3 However, Sec. 3.3 does not provide an explanation for composing these two operators. Namely, the whitening operator is estimated using normalized gradients, rather than the raw gradient. We will leave the investigation of operator composition for future work.
Adapprox Adapprox (Zhao et al. 2024b) uses low-rank approximation for Adam's 's second moments through randomized(SVD to boost the memory efficiency. However, its performance will be similar to Adam. In fact, Theorem 4.1 of |ACC proves that the converged & and represents the right and left singular vectors, coinciding with the rank-1 reconstruction. However, the proposed RACS] is different from rank-1 1 Adapprox in three main aspects: (1) Adapprox proposes to use low-rank EMA tracking on G2, whereas we use EMA directly on scaling vectors s, q. The resulting vectors are no longer singular vectors; (2) we do not have separate scaling apart from norm-grov growth limiter and user-defined scale, whereas Adapprox uses the scaling from Adafactor; (RACS do not use any first moment. Specifically, (1) is an important difference, since for any low-rank reconstruction with rank r > 1, is not guaranteed to positive, causing numerical issue during square-root inverse scaling. Adapprox adds a manually defined offset to ensure positivity. On the other hand, RACS is guaranteed to have positive s, q at each step from Cerron-Frobenice theorem. Therefore, RACS is numerically stable.
Adafactor Adafactor (Saazeer & Stern 2018) is another optimizer that uses rank-1 approximation to Adam's second moment. However, their scaling is derived by minimizing different norm, compared to Frobenius norm in this paper. Adapprox Zhao et al. (2024b) has demonstrated the advantages of using Frobenius norm compared to Adafactor with slightly better performance on LLM training.
RACS v.s. Adafactor The update rule of RACS resembles the Adafactor. Indeed, from Eq. (I). when one initializes S and Q as identity and performs one-step fixed point procedure, the resulting scaling matrix is identical to Adafactor. Intuitively, it seems to be surprising if one regards RACS as "multiple step" version of Adafactor and the performance difference comes solely from running the 1-step or multi-step procedure. Here, we show that extending from 1-step to multi-step can lead to drastically different parameter updates when combining with EMA of S and Q. We train 60M LLaMA model with RACS) and simultaneously calculate the Adafactor parameter update, followed by computing their relative difference as:
Rel. Diff
1 N
A
RACS
A Ad..fllaeorlll
(48)
|Ar,Adafactor// ||
where A, indicates the update for parameter i. Fig. 4 shows the relative difference curve throughout the training. We can observe the relative difference is small, but quickly increases as the training
41

---

Published as a conference paper at ICLR 2026
Average riaiive update diffenence from PACS and Adalaactor
No: Dfficeeee
2309
5000
1999
2900 10699
12993
1900
11593
51895
Figure 4: The relative update difference between RACS and Adafactor for training 60M LLaMA RACS
progresses. In the end, it settles at 0.2, representing RACS differs from Adafactor by 20% at each update step. This confirms that the RACS results in different update rule compared to Adafactor.
AdaDiag and one-sided SOAP We acknowledge that there exists two concurrent work: AdaDiag (Nguyen et al. 2025) and one-sided SOAP (Vyas et al., 2024), that are mathematically equivalent to OSOAP They are derived based on distinct view points. AdaDiag is based on the intuition to transform the gradient covariance matrix so that it is diagonalizable, where this diagonal matrix is estimated using Adam's 's second moments. One-sided SOAP, a memory-efficient version of SOAP, is proposed to based intuitions that only one-sided eigenspace is enough for LLM training. On the other hand, (OSOAPis is derived on the basis of the structured FIM view point. providing deeper connections to optimizers considered in this paper.
Apollo There is one concurrent work, Apollo (Zhu et al. 2024). that is also based on scaling the raw stochastic gradients for memory-efficient LLM training. It proposed to scale columns OR rows of the G' through scaling matrix estimated by similar procedure as Fira Chen et al. (2024a). The main idea is that column or row norm after scaling matches the column or row norm of the gradient from GaLore update. Thus, they require GaLore procedure to compute this update at each step. Our proposed RACS scales both columns and rows at the same time. The main differences are: (1) the scaling estimation in |RACS is different from Apollo; (2) the scaling scheme of RACSS is inspired by the generalization of normalization, providing theoretical support. They enjoy similar memory consumption (i.e. mn 2n + of Apollo compared mn + m + n++1 RACS).
improvement towards GaLore, which modifies the low-rank update to full-rank one by adding a compensation term. Their idea is similar to the compensation used in Alice The main differences is that proposed compensation is inspired by approximating FIM and has the theoretical foundation on its optimality, Also, the compensation strategy is different Fira. We have conduced the ablation study in Sec. ooto show the advantage of our approach.
Fira Fira Chen al. (2024a) was proposed
GaLore Comparing the Allcee procedure to GaLore (Algorithm we can clearly see that GaLore is special case of bl disabling tracking, switching and compensation. These three steps flows naturally from the structural approximation view point. This shows the advantage of this view point compared to ad-hoc choices when designing new efficient optimizers.
SARA Sara (Zhang et al. 2025) is concurrent work that aims to improve the quality of the projection by using importance sampling of singular vectors. This shares the similar intuition with our subspace switching mechanism (Algorithm IT, although the techniques differ: Sara relies importance sampling of singular vectors, while Alice mixes random sampling of less important directions with leading directions. We want emphasize that subspace switching is only one of
42

---

Published as a conference paper at ICLR 2026
the three components in Alice Based the ablation study in Table 4 the theoretically optimal compensation step offers the largest improvement in performance.
LDAdam LDAdam (Robert et al. 2024) is a recently proposed low-rank optimizer based on Adam with error-feedback mechanism. The goal of error-feedback is similar to the compensation step in Alice aiming at reducing the information loss from the low-rank projection. However, they tackle this problem different manner. Error-feedback mechanism modifies the gradient by adding back the information loss from the previous steps. However, LDAdam update still relies on the low-rank projection, resulting in low-rank update the end. Additionally, there no theoretical guarantees on the optimality of error-feedback mechanism. On the contrary, [llce leverages the optimal compensation step, which directly add the information loss back. Therefore, the parameter update at each step of Alice is guaranteed to be full-rank.
1.6 DISCUSSION OF LOW-RANK EXTENSION FRAMEWORK
First, we derive how to decompose the full-rank update of OSOAP into low-rank update and its residuals in the main text. We assume Ú -U,U,, U,], where U. U, are defined in Sec. 5.3
ÚTG
Mat Mat(F- (--2 -ū
|[(ÚTG]]]
G)
U UT
G
-UU.Udl
U"
G
UTG
UTG
-UU
U ))
UAG
E
UTG
UUT
=UU
U.G
VEKUTG
UTG
U
+U.
[UTG
Moore-Penrood inverse In Eq. 13 we define F- using the pseudo-inverse since each block in F. low-rank and invertible. To be precise, for any block U.D..U!. its pseudo-inverse be easily verified U.D_!U² by checking the definition. Similar typical inverse, for any block diagonal F.,, the pseudo-inverse equivalent applying pseudo-ing inverse of each blocks. The square-root pseudo inverse can be defined through similar way.
Similarly, for the proposed compensation, we can easily verify its vectorized format (ignoring the subscript for simiplicity)
Vee(U.U GS) -S @ U. g
=(S U.U!) ag
where the - is the eseudo-inverse the above. The second inequality can be easily verified i
by the following facts: (1) the square root pseudo inverse of U,U is itself; (2) the square root pseudo-inverse of block-diagonal matrix is the pseudo-inverse of each individual block.
43

---

Published as a conference paper at ICLR 2026
J EXPERIMENT DETAILS
In this section, will include the detailed setup, hyperparameters and additional experiment results.
K IMPLEMENTATION DETAILS OF BASELINES
GaLore We leveraged the official GaLore package released by the author (attps://githuz com/ Jlawelzzhao/GaLores
Fira Since the main difference compared to GaLore is the additional compensation term, we follow the implementation of the official Fira repo chttpss://ithubb htw.:u...omxxxiichhcccc xichen caiii-y/Friarr and add the compensation GaLore.
K.1 EXPERIMENT SETUP FOR PRETRAINING LLAMA
We context length of 256 and batch size of 128 with gradient accumulations, apart from 60M and 1.3B (batch size 256 with 2 gradient accumulations). RACS Alice are applied to all linear modules of attention and MLPs, others optimized with Adam. We use the first 10% of the total training steps as the warm-up, followed by a cosine learning rate decay to 10% of the original learning rate. For hyperparameters,
perform grid search our proposed RACS Alice well GoLore, Fira and full-rank Adam. For other methods, we directly cite their results in (Zhu et al. 2024). For Adam, we use 0.9 and 0.999 for $1 and B2, and enable the bias correction without weight decay. Table O summarizes the hyperparameters used for both GaLore and Fira. Table 7 summarizes the hyperparameters for Adam optimizer. Table X and Table TO TT summarizes the hyperparameters used for RACS AAice respectively. In addition, for all methods with Fira limiter, we use threshold
y = 1.01 as suggested in Chen et al. (2024a). For we use Adam to train the last layer of LLaMA, following the same setup Apollo for fair comparison. In summary, Alice GaLore and
Fira do not use Adam to train the last layer to fully test the capabilities of low-rank methods, whereas Apollo, [RACS] and Adam utilize Adam for last layer. The total training steps for 60M, 130M, 350M and 1.3B are 10K, 20K, 60K and 100K, respectively. These correspond to 1.1B, 2.6B, 7.8B and 13. 1B training tokens, summarized in Table All experiments conduced NVIDIA A100 GPUs.
Table 7: The hyperparameters used for Adam
Table 6: The hyperparameters for GaLore and
optimizer.
Fira
learning rate
update scale
rank
update
learning rate
correct bias
60M
0.02
0,3
128
200
60M
0.001
0.9
0.999
True
130M
0.02
256
200
130M
0.001
0.9
0.999
True
350M
0.02
0.3
256
200
350M
0.001
0.999
True
1.30
0.01
0.25
512
200
1.3B
7 10
0.9
0.999
True
Table 8: The hyperparameters
Table 9: Model architectures and training steps
for RACS)
Hidden
Intermediate
Heads
Layers
Steps
Data
learning rate
8
scale
60M
512
1376
8
8
10K
1.3B
60M
0.02
0.9
0.05
130M
768
2048
12
12
20K
2.6B
130M
0.02
0.9
0.05
350M
0.02
0.9
0.05
350M
60K
7.8B
1.3B
0.02
0.02
1.3B
4096
5461
24
32
100K
13.1B
Table 10: The hyperparmeters for Alice optimizer
learning
scale
compensation scale
15g
updale interval K
rank
leading basis number
60M
0.02
0.3
0.4
0.9
0.9
0.999
200
128
40
0.02
0.999
256
40
350M
0.02
0.3
0.4
0.9
0.9
0.999
200
256
40
1.3B
0.02
0.25
0.2
0.9
0.9
0.999
512
160
44

---

Published as a conference paper at ICLR 2026
K.2 SETUP OF 1B V.S. 7B
We mainly follow the setup as (Zuu et al. 2024). For IB model, we use 16 per-device batch size with 8xA100s and gradient accumulations, which 512 batch size in total, same as 7B model. For Alicel we use learning rate 0.02, scale a = 0.3 and compensation scale a. = 0.2 with rank 512, leading basis number = 160. We also use full-rank Adam train the last layer for better performance. For memory estimation, we assume the use of BF16 format. For 8-bit optimizers, we model.
assume weights are stored BF16, but optimizer states use FP8. GaLore uses r 1024 for 7B
K.3 MEMORY ESTIMATION
Following the setup of Zhao al. 2024a), we provide the estimated GPU memory for each optimizers due to the difficulty of directly measuring their practical memory consumption without considering the activations, internal storage for gradient, etc. We assume BF16 format, which consuming Bytes per element. For example, 1273M parameters are optimized by [llic and 65M parameters are optimized by Adam for 1B model with rank 512. Therefore, [Alice part will consume 3.86GB and Adam will consume 0.37GB, summing to 4.42GB for Alice optimizer.
We also report the actual memory footprint with BF16 format. We token batch size of 25, following the same setup as /hao et al.] (2024a). "-layerwise" represents layer-wise training where we only store the gradient of the current back-propagated layer.
K.4 THROUGHPUT ESTIMATION
We report both the actual throughput and effective throughput. The effective throughput of target optimizer compared to reference optimizer is defined as the ratio of training tokens consumed from the target optimizer w.r.t total time consumption of reference optimizer when reaching the same evaluation loss of the reference optimizer at the end of training. Compared to the standard throughput, this considers the speed-up effect.
K.5 ADDITIONAL PRETRAIN RESULTS
Fig. S presents additional training curves with 60M, 130M and 350M models sizes. For all cases, the proposed Alice and RACS outperforms the baselines with noticable margin, and achieves clear speed-ups compared to full-rank Adam.
K.6 RESULTS FOR ABLATION STUDIES
For all ablations, we consider 130M I.LMAA model.
Sctup: ablation of tracking We disable the compensation step, and verify the effect of low-rank tracking under our proposed switch strategy or purely replying on EVD The other hyperparameters follow the pre-training setup as Sec. K.I
Sctup: switch strategy For this setup, we enable low-rank tracking. Apart from Gaussian, 40 as the leading basis number. The Gaussian distribution is a standard isotropic Gaussian with zero mean and unit variance.
Setup: compensation strategy We enable the tracking and switching, but with different compensa- tion terms. For Fira, we directly leverage the compensation proposed in Chen et al. (2024a). Firat modifies original Fira by the following two steps: (1) rescale the Fira compensation have the same l2 norm as the Alice low-rank update; (2) multiply this compensation by a scale parameter like az in Algorithm
We found this empirical trick boosts the performance of Fira.
Effects of last layer One crucial setup difference during evaluation for low-rank methods is whether the last layer is trained by full-rank Adam or not. Most previous work train the last layer, arguably one of the most important layer (Zhao et al. 2024c), using full-rank Adam. This effectively reduces
45

---

Published as a conference paper at ICLR 2026
Evall PPL with model size: 60m
Eval PPL with model size: 130m
30.0
47.
#ice
Bra
Galane
Galare
Alos- houd
40.0
27.5
Evel
271855
11.0
2.00
258
24-
2.08
22
12000
14828
10000
Update Stop
Update Stop
(a)
(b)
Eval PPL with model size: 350m
Cellans
Acc im head
22:11
Evall
201
Updale Stup
Figure 5: Additional LLaMA C4 pretrain performance curve. (a), (b) and (c) represents the 60M, 130M and 350M, respectively. "+lm head" represents that the last layer of LLaMA is trained by full-rank Adam.
46

---

Published as a conference paper at ICLR 2026
Through put Method and Model Size
100000
Adam
Fira
RACS
Alkc
Throughout
60000
40000
20000
(a) Absolute throughput
Effective Throughput by Method and Model Size
Adam
Ara
175000
RACS
Alice
150000
Elfeciee thrughret
100000
75000
23000
60m
1.0m
350m
3b
Model Size
(b) Effective throughput
Figure Throughput of various methods. (a) this reports the absolute throughput, representing the number of training token processed per second. (b) the effective throughput using Adam as the reference optimizer. This represents the absolute throughput adjusted by the speed-up factor. The effective throughput of GaLore is 0 since it under-performs the Adam.
47

---

Published as a conference paper at ICLR 2026
Memory footprint of d optimizers and model size
Galore
Galore-layrrvise
Fiaa
Rra layerwise
40
RACS layerwise
Alice-0
Merrory Protprinn (68)
0laayerisee
Alice
30
10
350m
1b
76
Model Size
Figure 7: The memory footprint of various optimizers. We use token batch size of 1 following the same setup as Zhhao al. (2024a) under BF16 format. The suffix "layerwise" represents the memory consumption when enabling layerwise training so that only the gradient of the current layer is stored.
the performance gap compared to full-rank method, and does not reveal their true capabilities. We investigate the effect of the last layer to A compared to GaLore. From the Table 2 we can see that for all model sizes, GaLore and Fira are greatly affected by this effect. Training last layer with full-rank Adam will boost their performance significantly. On the other hand. Aliic is less impacted with marginally worse performance. For example, Fig. 8(d) shows the training curve comparison with I 130M model. When the rank is sufficiently large (i.c. 60M model), Fig. (hh)shows shows that using Adam to train the last layer even decreases the performance. These serve as evidences that Alice is better optimizer than GaLore and Fira.
Effect EMA in RACS and q. We investigate the importance of [MA scheme. From Fig. 8(e) RACS without [MA performs much worse than RACS suggesting the EMA is necessary for satisfactory performance of RACS
The only internal states inside RACS are the EMA tracking of two vectors
Hyperparameter sensitivity analysis: Adam We consider 130M Llama model, and the other setup follows the same setting the pretraining experiments. For Adam, consider range of learning t
rate from 3 x 10-4 to 3 x 10-3, 31 from (0.9, 0.8, 0.6, 0.4] and 32 from [0.999, 0.95, 0.8, 0.6, 0.4]. The total budget 120 runs.
For 130M model, we found that the best performing hyperparameters for Adam are learning rate 0.001 with betas 3 - 0.8, 32 0.95. This achieves ppl of 22.89. Surprisingly, apart from RACS (ppl 22.67) and Alice (pp! 21.95). fine-tuned Adam outperforms or is on par with all other candidates. despite some of them claims achieve better convergence than Adam. One hypothesis is that Adam is sensitive to beta parameter, and most existing literature may use the untuned beta (i.c. B1 = 0.9, 32 = 0.999) for comparison. Even compared to fine-tuned Adam, Aliiee still achieves 1.4r speed-up. [I reports a Adam performance with selected set of hyperparameters. We observe that Adam's 's performance can drop significantly even by slightly changing the 32, For example, with 31 = 0.8. changing B2 from 0.95 to 0.8 decreases the ppl from 22.89 to 24.51. Due to the expensive nature of each run, fine-tuned Adam not practical.
Hy Elyperaarameter sensitivity analysis: Alice We follow the same setup the sensitivity analysis of Adam. We consider the following hyperparameters Alice (1) compensation scale a from 0.1 0.9; (2) update interval K from [50, 100, 200, 300, 400]; (3) rank from (16, (2, 64, 128, 256):
48

---

Published as a conference paper at ICLR 2026
Effects low-rank tracking
Effects switch strategy
tracking trnding
bunk
conatmmnn todd
253007
Update Step
Updato Step
(a) Effect of tracking
(b) Effect of switching
Effects compensation
Effects EMA
1112
#nd tiineed
Galorre
25.1144
27:8827
22.4553
55.2854
22,2952
400
200g
Step
0003
1550
400
Update Step
1003
1000
10000
1823
1150
(c) Effect of compensation
(d) Effect of last layer. "+Im head" represents the last. layer is trained by full-rank Adam.
Effects of EMA scheme
cval po)
pdatt
200
1853
6147
(e)
Effect EMA in RACS
Figure& 8: The pre-training curve to verify the effectiveness of the design choice. We consider 130M model size.
Betas
(0440.6 0.6)
(0.4, 0.95)
(0.6 .0.8)
(0.6, 0.95)
(0.8 .0.8)
(0.8, 0.95)
(0.9.0.8)
(0.9.0.95)
Ppl
25.18
23.44
23.79
23.05
24.51
22,89
24.90
23,36
Table 11: The validation ppl. of Adam optimizer with different choice of betas. Due to the large num- ber of runs (120 in total), only report Adam with Ir 0.001 and some selected betas combinations. This is enough to prove that Adam is sensitive to hyperparameter selection.
49

---

Published as a conference paper at ICLR 2026
Figen-Space Cosine Similarity (Index
Bigen ce Casine Simil Undex 2)
1.0
1.0
0.8
0.8
Simlarty
Costre Sintarty
tracking smrtch tracking switch
0.0
2500
5000
7520
1,0000
12500
15000
17500
20000
2500
5000
7500
10000
125.0
15000
17500
20000
Stoplindax DataFrame
Stapsindox DataFromz
(a) Index
(b) Index 2
Eigen-Space Similarity
Eigen Space Similarity
Siilartty
trocking ho switch i accking no switch
PO tracdiing no Se tth tracking king no switch
(c) Index
(d) Index 8
Figure 9: The cosine similarity between eigenvectors per 200 steps. Since the update interval of subspace is 200 steps, this essentially compare the similarity before and after updating the projection U with a certain index. We always arrange the eigenvectors in a descending order based on the eigenvalues.
50

---

Published as a conference paper at ICLR 2026
and (4) leading basis number from [5, 10, 20, 40, 80, 160]. Due to the combinatorial nature, only change one particular hyperparameter Afice at a time and fix the others as Table I0
Tables 12 15 shows the performance of Alice under different choice of hyperparameters. We observed that Alice is not sensitive to the choice of hyperparameters. Among those runs (25 in total), 23 of them outperform fine-tuned Adam, proving that the advantage [f/Alceccome from its principled design, rather than carefully selected hyperparameters.
Compensation scale
0.1
0.2
0.3
0.4
0.5
0.6
0.7
0.8
0.9
Ppl.
23.04
22.29
22.08
21.95
22.02
22.05
22.12
12.12
22.19
22.27
Table 12: The performance of Alice with different compensation scale de-
Update interval K
50
100
200
300
400
Ppl.
23.01
22.09
21.95
21.99
22.09
Table 13:
The performance of Alice with different choice of update interval K.
Rank r Ppl. 22.20
16
32
64
128
256
22.03
22.00
22.07
21.95
Table 14: The performance of Alice with different choice of rank T.
Leading basis number
40
80
160
Ppl.
22.53
22.33
22.21
21.95
21.95
21.90
Table 15: The performance of Alice with different choice of leading basis number 1.
K.7 GLUE FINETUNE PERFORMANCE
Although the primary goal of this paper is to design the efficient optimizer for pre-training LLMs, we also conduct the preliminary finetune experiments with Alice using GLUE (Wang 2018).
Setup We follow the
setup as (Zhao et al., 2024a) with RoBERTa Liu 2019) model. We select the rank r = for both GaLore and Alice! Tables [6 and 17 reports the hyperparameters for GaLore and Alice By default, report prediction accuracy the metric for most of the dataset. For CoLA, we use Matthews Coattess_Corrllttonn, and for STS-B, we report Pearson Correlation.
Table 18 18 reports the GLUE finetune performance of Adam and Alice is clear that Alice consiitenty outperforms the GaLore across datasets. This is expected since each update step of Alice is full-rank compared to the low-rank nature of GaLore. Also, GaLore is special case of Alice
MNLI
SST-2
MRPC
CoLA
QNLI
QQP
RTE
STS-B
Batch size Epochs
16
16
16
32
16
16
16
16
30
30
30
30
30
30
30
Learning rate
le-5
2e-5
2e-5
le-5
le-5
2e-5
2e-5
3e-5
Scale a
Sequence length
512
Update interval
500
Table 16: Galore hyperparameters for GLUE finetune experiments.
51

---

Published as a conference paper at ICLR 2026
MNLI
SST-2
MRPC
CoLA
QNLI
QQP
RTE
STS-B
Batch size
16
16
16
32
16
16
16
16
Epochs
30
30
30
30
30
30
30
30
Learning rate
1e-5
1e-5
2e-5
2e-5
le-5
1e-5
le-5
2e-5
Scale a
4
Sequence length
512
Update interval
500
Leading basis number
8
3
6
6
3
6
Compensation scale
0.1
0.07
0.1
0.05
0.1
0.1
0.3
0.3
Table 17:
Alice hyperparameters for GLUE finetune experiments.
MNLI
SST-2
MRPC
CoLA
QNLI
QQP
RTE
STS-B
Mean
GaLore Alice
87.31
94.72
89.71
60.34
92.75
91.13
79.78
91.11
85.86
87.45
95.18
90.69
65.04
92.90
91.45
80.51
91.43
86.83
Table 18: The GLUE finetune performance of Adam and Alice
52


Reviews:
[{"review_id": "3QWyh6ELgo", "reviewer": "Reviewer_CJCW", "summary": "Memory, sample efficiency, and throughput are key performance metrics for optimization algorithms, and algorithmic development is essential for mitigating the cost of LLM training. This work leverages an insightful lens of Fisher Information Matrix (FIM) preconditioning (Natural Gradient Descent view) to study successful optimizers such as Adam, Shampoo, Muon, and Soap, and unifies them all under this lens by studying the specific structure of FIM approximation (different variants of Kronecker product and block diagonal structure) that each of these optimizers is made of. Interestingly, the update rule of the one-sided Soap algorithm is also derived and analyzed by leveraging block-diagonal structure, and together the one-sided Soap update rule plus the derived structure is referred to as Generalized Adam (GAdam).\n\nInspired by FIM unification, this work hence argues that exploration of novel structures for FIM approximations is the key design principle for systematically developing new optimizers to further balance the optimizers' performance metrics, especially memory vs sample efficiency. Building on this design principle, the paper proposes two novel optimization algorithms. The first proposed algorithm is called RACS, where the preconditioner is structured as a Kronecker product of two diagonal matrices, and by minimizing the FIM approximation error (measured in Frobenius norm), the two diagonal matrices are derived subsequently. This optimizer hence reduces the second moment memory significantly and matches the memory complexity of AdaFactor (for instance, Adam's second moment needs O(mn) while RACS needs O(m+n)).\n\nMoreover, this work also proposes another novel low-rank optimizer called ALICE, achieving GaLore-like memory while significantly improving both sample efficiency and throughput. ALICE extends the GAdam algorithm (one-sided Soap update rule) to low-rank subspace optimizer settings by maintaining only the top r eigenvectors of the full-rank projection in GAdam, and then the projected gradients are accumulated in the subspace akin to GaLore. Subspace switching and compensation methods are also proposed to account for the residual update term that remains orthogonal to the projection matrix and is wiped out during projection.\n\nThe two proposed algorithms are validated on LLaMA pre-training and GLUE benchmarks compared to AdamW, GaLore, Fira, Apollo, and Muon. The reported results indicate that these memory-efficient optimizers (ALICE and RACS) not only match the final validation perplexity of the best baseline but both improve the wall-clock time by a factor of 2 compared to Adam.", "questions": "1. Could the authors please address my critique regarding RACS, and particularly explain the key algorithmic differences between RACS and AdaFactor beyond performing more steps of power iterations? This is indeed a very important concern to address, as AdaFactor does not achieve anywhere close to the RACS speedup compared to AdamW, and having more power iterations as the only key algorithmic component does not seem sufficient to explain these gains.\n\n2. Could the authors please provide a detailed comparison of the compensation step in ALICE to LDAdam? Line 467 mentions Fira as the closest work for compensation, but LDAdam has also made significant contributions regarding this compensation mechanism, which merits discussion and comparison.\n\n3. Proposition 4.1 in this work appears to be a diagonal version of Proposition 1 (Section 3.1 in [3]). It would be interesting to see what additional challenges the proof of Proposition 4.1 introduces compared to Proposition 1 in [3], and how these relate to power iteration convergence.\n\n\n\n[1] Wen, Kaiyue, et al. \"Fantastic pretraining optimizers and where to find them.\" arXiv preprint arXiv:2509.02046 (2025).\n\n[2] Robert, Thomas, et al. \"Ldadam: Adaptive optimization from low-dimensional gradient statistics.\" arXiv preprint arXiv:2410.16103 (2024).\n\n[3] Morwani, Depen, et al. \"A New Perspective on Shampoo's Preconditioner.\" arXiv preprint arXiv:2406.17748 (2024).\n\n[4] Shazeer, Noam, and Mitchell Stern. \"Adafactor: Adaptive learning rates with sublinear memory cost.\" International Conference on Machine Learning. PMLR, 2018.", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well-presented and clear. The proposed algorithms are derived based on rigorous theoretical principles, and experiments are comprehensive and profile all key performance metrics (throughput, memory, and perplexity) across the LLaMA model family (from 60M to 1.3B). Comprehensive ablations verify the effectiveness of the key algorithmic components introduced for the two proposed optimizers.\n\n2. This paper makes it clear that the most adopted and successful optimizers in deep learning indeed are different choices of FIM approximation. Principled design of optimizers based on exploring the space of FIM approximation structures is a very clever and strong idea. Finding principles for one-sided Soap based on this framework is an interesting contribution.\n\n3. Both proposed optimizers, ALICE and RACS, achieve consistent speedup by a factor of 2 (measured in wall-clock time) and significant memory savings compared to AdamW. These gains represent very strong results, and upon reproduction of these results, the ALICE and RACS algorithms have the potential to be adopted as mainstream optimizers for LLM training based on reported results. Especially considering the benchmark of current state-of-the-art optimizers in similar pre-training settings [1], the speedup of the best SOTA optimizer compared to Adam is around 1.1-1.4, and Muon is among the best-performing optimizers. Interestingly, the paper reports both ALICE and RACS outperform Muon significantly while using less memory.", "weaknesses": "### Critiques on the FIM Framework (Section 3)\n\nApplying the Shampoo preconditioner directly on gradients can recover spectral normalization (whitening) as well, and hence whitening can be understood without leveraging Proposition 3.3. Moreover, [3] already finds the optimal solution to Equation 2 (line 147), where the optimal R and L are computed as principal singular vectors of reshaped F. Then [3] leverages power iteration for approximating the optimal solution of Equation 2, and hence minimizing the proposed upper bound in Theorem 3.2 does not provide significant additional insight when compared to established results in [3]. For Theorem 3.4, if we assume all D_i are equal, then the structure would be equal to the whitening structure that is used to derive Equation 3, and hence the derivation of Equation 6 appears redundant. Moreover, since GAdam appears to be equivalent to one-sided Soap, and this algorithm is actually discussed in the original Soap paper, the justification for not referring to GAdam as one-sided Soap is unclear. While the derived principle has novelty, the update rule remains the same.\n\n### Novelty of RACS Compared to AdaFactor\n\nIf we let the initial S and Q be identity matrices with proper scaling, then one iteration of the fixed point procedure in Equation 11 would indeed recover the AdaFactor update. Moreover, when applying the structure in Equation 10 to gradients as in Equation 1, we can write:\n$\\text{Vec}(\\Delta W) = -\\lambda (S \\otimes Q)^{-1/2} \\text{Vec}(G)$. Since both S and Q are diagonal, $S \\otimes Q$ would be diagonal as well, and hence we can write $\\Delta W = -\\lambda \\frac{G}{\\sqrt{V}}$, where $V = qs^\\top$. Since Proposition 4.1 shows that q and s converge to the left and right singular vectors of $G^2$ up to scaling factors, this shows that V would be a rank-one SVD approximation of $G^2$, and the iterative fixed point procedure of RACS is essentially a power iteration that approximates those principal vectors. Moreover, the original AdaFactor (Algorithm 4 in [4]) similarly maintains momentum on q and s. Hence, RACS appears to be very similar to AdaFactor. RACS performs 5 iterations of Equation 11, while AdaFactor performs one iteration. Therefore, the novelty of the RACS update rule could be considered limited, as RACS appears to be another rank-one approximation method of the second moment, which AdaFactor has already explored and studied in detail.\n\n### Reproducibility\n\nThe code has not been made available, and considering the significant speedup reported in this work and the potential for improving upon the best existing optimizers in the field, making the source code and implementation public would greatly benefit the research community and enable better evaluation of this paper. Even sharing the source code for the 130M setup would be highly valuable.", "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": "Me8atQqtIA", "reviewer": "Reviewer_mnjC", "summary": "In this paper, the authors have proposed a unified framework for both memory and computational-efficient optimizers for LLMs. To achieve this goal, the authors have used the structured Fisher information matrix (FIM) approximation. Many existing optimizers, such as Adam and Shampoo, are approximations to the gradient descent update. FIM, on the other hand, is replaced by certain structured forms, such as block diagonal and Kronecker product matrices.", "questions": "Galore has numerous follow-up works, such as Galore 2 [1], Golore [2], and Sara [3]. Could the authors give a comparison with these works? \n\n\n[1] DiJia Su, Andrew Gu, Jane Xu, Yuandong Tian, and Jiawei Zhao. \"Galore 2: Large-scale llm pre-training by gradient low-rank projection.\" arXiv preprint arXiv:2504.20437 (2025).\n\n[2] Yutong He, Pengrui Li, Yipeng Hu, Chuyan Chen, and Kun Yuan. \"Subspace optimization for large language models with convergence guarantees.\" ICML'25.\n\n[3] Haochen Zhang, Junze Yin, Guanchu Wang, Zirui Liu, Tianyi Zhang, Anshumali Shrivastava, Lin Yang, and Vladimir Braverman. \"Breaking the Frozen Subspace: Importance Sampling for Low-Rank Optimization in LLM Pretraining\". NeurIPS'25.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "The strengths of this paper are summarized as follows:\n\n1. This paper has a novel optimizer design. Row and Column Scaled SGD (RACS) introduces a minimal overhead two sided scaling method. It achieves strong results with almost no additional memory. \n\n2. Empirically, the authors have shown experimental results on different model scales, and they are compared with many baselines, such as the well-known Adam and Galore. The experimental results have shown it has consistent performance improvements.\n\n3. Theoretically, the authors have proved the convergence of Alice under the continuous-time setup.", "weaknesses": "The weaknesses of this paper are summarized as follows:\n\n1. Theoretically, the authors have given the convergence analysis only in the continuous time setting and “leave the general treatment of the convergence to future work”. It would be better if the authors could give a brief discussion about what else needs to be done, what the difficulties are, and what the differences are with the continuous-time setup.\n\n2. Empirically, the authors only analyze “pre-training LLaMA on the C4 dataset”. It would be better to include more evaluations in a more diverse domain. Also, the model sizes are small. Nowadays, people are more interested in larger LLMs. It would be better to include the experimental results on larger models.", "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": "7CTtltrE0Q", "reviewer": "Reviewer_phAX", "summary": "This paper proposes a unifying framework for designing memory-efficient LLM optimizers through structured Fisher Information Matrix (FIM) approximation. The authors show that existing optimizers (Adam, Shampoo, gradient normalization, etc.) can be viewed as special cases of FIM approximation with different structural constraints, and propose two design principles: (1) selecting structures that balance generality and efficiency, and (2) applying a low-rank extension framework. Based on these principles, they introduce two novel optimizers—RACS (achieving SGD-like memory with row-column scaling) and Alice (a low-rank variant of generalized Adam with tracking, switching, and compensation mechanisms)—which demonstrate over 2× faster convergence than Adam on LLaMA pre-training tasks while maintaining low memory overhead.", "questions": "Q1. Can you provide convergence analysis for Alice under discrete-time dynamics, including convergence rates in terms of learning rate η, rank r, and switching frequency K, and how the one-step refinement approximation error affects final convergence?\n\nQ2. Can you provide detailed profiling showing time spent on EVD approximation, subspace switching, and compensation computation, and demonstrate how throughput scales with update interval K to help practitioners optimize the 15% throughput penalty?\n\nQ3. Why is random uniform sampling from the complement space optimal in Algorithm 1, and have you compared against informed strategies like gradient-based importance sampling or residual-guided selection using Σ_t from Theorem G.8?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Originality: The unifying FIM approximation framework is highly original, elegantly connecting diverse optimizers (Adam, Shampoo, normalization, Muon) under a principled theoretical lens. The low-rank extension framework with subspace switching and optimal compensation represents creative innovation beyond existing ad-hoc approaches. The insight linking GaLore to Gadam provides valuable new perspective on existing methods. \n\n2. Quality: The mathematical formulations are rigorous, with analytical solutions derived for various structural constraints. Experiments are comprehensive and well-designed, spanning multiple model sizes (60M-1.3B) and demonstrating consistent, substantial improvements. The convergence analysis under continuous-time setup adds theoretical grounding.\n\n3. Clarity: The paper is well-structured with logical progression from framework exposition to practical optimizer design. The step-by-step derivation of existing optimizers as special cases effectively illustrates the framework's unifying power. The two design principles provide clear, actionable guidance.\n\n4. Significance: High practical impact—Alice achieves over 2× convergence speedup compared to Adam with minimal memory overhead, addressing critical LLM training efficiency challenges. The framework provides valuable conceptual foundation for systematic optimizer design, moving the field beyond heuristic approaches toward principled methodology. The work opens promising research directions for incorporating model structure and task-specific information into optimizer design.", "weaknesses": "1. The convergence analysis is limited to continuous-time setup (Theorem H.2), which is a gap given that practical optimizers operate in discrete time. The one-step refinement approximation in Theorem 3.4 for Gadam lacks justification for why this approximation is reasonable or how the error propagates.\n\n2. All pretraining uses only LLaMA on C4. Section K.7 adds GLUE on RoBERTa, but this is insufficient. Generalization to other architectures (GPT, T5, Mamba) and domains (code, multilingual) is unclear.\n\n3. Table 2 shows Alice has 15% throughput drop on 1B models (45,523 vs 53,588 tokens/sec), but lacks breakdown of where this overhead originates. Specifically, no profiling data separates the cost of: (1) EVD approximation via subspace iteration (O(m²r) every K steps), (2) subspace switching with QR decomposition, and (3) per-step compensation computation. Without this profiling, practitioners cannot determine which components to optimize (e.g., increasing K to amortize EVD cost, or simplifying Eq. 14 to reduce compensation overhead) or assess when the 15% penalty is acceptable given 2× step-wise speedup.", "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": "J1M8is5uje", "reviewer": "Reviewer_1U1o", "summary": "The paper advocates a structured Fisher Information Matrix (FIM) approximation as a unifying lens for designing efficient optimizers for LLM training. It shows that a range of optimizers/gradient operators—including Adam, Shampoo, normalization/whitening, and SOAP/AdaDiag++—arise as solutions to a Frobenius-norm projection of the empirical FIM onto different structure families. On LLaMA pretraining, RACS/Alice outperform Adam and strong memory-efficient baselines; Alice achieves >2× step-speedup over Adam with comparable memory to GaLore, and RACS delivers strong results with ~SGD-like memory. The contribution is mainly a unifying viewpoint and design methodology, not a definitive theory of LLM optimization.", "questions": "see weaknesses", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "1. This paper provides unfied and clear perspective: casting popular methods as structured FIM projections.\n\n2. On C4 pretraining, Alice yields 2.22×–2.82× step-speedup vs. Adam with strong effective throughput; RACS beats multiple memory-efficient baselines (GaLore, Fira, Apollo, Muon) under comparable or lower memory. \n\n3. The writing is good and easy to follow.\n\n4. Solid theorem analysis is provided.\n\n5. Both pretraining and finetuning experiments are provided, showing the superiority of the proposed method under different cases.", "weaknesses": "Overall, I believe this is a good paper, some small cons:\n\n1. The font size in some tables and figures are too small for potential readers, for example, Figure 1 and Table 1. It would be better if the font size be larger.\n\n2. The authors may consider adding illustrative figures to accompany the complex mathematical formulations. Visual aids would help readers grasp the core ideas and intuitions behind the proposed methods more intuitively and easily\n\n3. The current experiments are all conducted on models with fewer than 1B parameters. It would be better to see if the proposed method can scale up to 7B model pre-training, as Galore and Fira do. Alternatively, if constrained by computational resources, the authors may consider offering more in-depth discussions or analyses to compensate for the lack of large-scale experiments.", "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)