4-bit Shampoo for Memory-Efficient Network Training

Authors: sike  wang
pan  zhou
jia  li
hua  huang
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-09 09:40:58

Abstract: Second-order optimizers, maintaining a matrix termed a preconditioner, are superior to first-order optimizers in both theory and practice.
The states forming the preconditioner and its inverse root restrict the maximum size of models trained by second-order optimizers. To address this, compressing 32-bit optimizer states to lower bitwidths has shown promise in reducing memory usage. However, current approaches only pertain to first-order optimizers. In this paper, we propose the first 4-bit second-order optimizers, exemplified by 4-bit Shampoo, maintaining performance similar to that of 32-bit ones. We show that quantizing the eigenvector matrix of the preconditioner in 4-bit Shampoo is remarkably better than quantizing the preconditioner itself both theoretically and experimentally. By rectifying the orthogonality of the quantized eigenvector matrix, we enhance the approximation of the preconditioner's eigenvector matrix, which also benefits the computation of its inverse 4-th root. Besides, we find that linear square quantization slightly outperforms dynamic tree quantization when quantizing second-order optimizer states. Evaluation on various networks for image classification and natural language modeling demonstrates that our 4-bit Shampoo achieves comparable performance to its 32-bit counterpart while being more memory-efficient.

Main paper: 
 introduction: 
 Deep neural networks (DNNs) have achieved great success in numerous fields, e.g., computer vision [20] , natural language processing [38] , and speech recognition [16] . A significant part of such success is attributed to first-order optimizers such as stochastic gradient descent with momentum (SGDM) [31] and AdamW [29] . Second-order optimizers, including K-FAC [30] , Shampoo [18] , AdaBK [41] , CASPR [13] , and Sophia [27] , show great convergence properties, but often involve noticeable computation and memory costs. Anil et al. [2] provided several practical techniques for second-order optimizers to achieve substantial wall-clock time improvements over traditional first-order optimizers. The fast convergence property of second-order optimizers benefits from preconditioning the gradient with a matrix known as a preconditioner. The optimizer states for constructing the preconditioner and its inverse root can speed up optimization compared to first-order optimizers, but consume memory that could be used for model parameters, limiting the maximum model size trained within a given memory budget. With the increase in model size, the memory utilized by optimizer states can become a predominant factor in memory usage. This is the primary obstacle hindering the widespread use of second-order optimizers in the era of large models.There are two main attempts to reduce memory consumed by optimizer states. Factorization uses lowrank approximation to optimizer states. This strategy has been applied to first-order optimizers [35, 3] (b) ViT-Base/32 on ImageNet-1k Figure 1 : Visualization of test accuracies and total GPU memory costs of vision transformers. 4-bit Shampoo (naive) quantizes the preconditioner, while 4-bit Shampoo (our) quantizes its eigenvector matrix.and second-order optimizers [14, 40] . In a comparable but distinct line of work, quantization utilizes low-bit to compress 32-bit optimizer states. Quantization is attractive due to its simplicity and wide applicability, which has been applied to first-order optimizers [8, 26] . Applying quantization to second-order optimizers poses a greater challenge, as first-order optimizers' states are elementwise, whereas second-order optimizers rely on matrix operations. To our knowledge, it has not been attempted before.In this paper, we present the first second-order optimizers with 4-bit optimizer states by taking Shampoo [18] as an example, while preserving the performance achieved with 32-bit optimizer states. While our focus is on Shampoo, we believe that our approach could also be applied to other second-order optimizers (see Table 4 ). Our main contributions are highlighted below.Firstly, to maintain 32-bit performance, we propose quantizing the eigenvector matrix of a preconditioner in 4-bit Shampoo, rather than the preconditioner itself. The reason is that the small singular values of the preconditioner matter. Directly quantizing the preconditioner via block-wise quantization [8] at 4-bit precision can significantly alter the small singular values, leading to a drastic change in its inverse 4-th root and thus harming 4-bit Shampoo's performance. Quantizing the eigenvector matrix can help alleviate this issue, which is supported by experimental validation and theoretical insight. Additionally, with the eigenvector matrix, computing the inverse 4-th root is straightforward, ensuring that quantizing the eigenvector matrix does not lead to a rise in the total wall-clock time compared to quantizing the preconditioner (see Figure 1 ). Secondly, we present two techniques for enhancing performance. As the eigenvector matrix of a preconditioner is orthogonal, we apply Björck orthonormalization [4] to rectify the orthogonality of the quantized eigenvector matrix, leading to improved approximation of preconditioner's eigenvector matrix and facilitating computation of its inverse 4-th root. Additionally, we observe that linear square quantization outperforms dynamic tree quantization [7] marginally when quantizing second-order optimizer states. The superiority of our developed 4-bit Shampoo is demonstrated in Figure 1 .Finally, we evaluate our 4-bit Shampoo on different image classification and natural language modeling tasks using convolutional neural network (CNN) and transformer architectures. Across all these benchmarks, our 4-bit Shampoo achieves similarly fast convergence comparable to its 32-bit counterpart, with no significant increase in losses for the trained models. Our 4-bit Shampoo uses less memory than its 32-bit counterpart, allowing for training of larger models with given resources. 
preliminaries: 
 In this section, we present Shampoo and its implementation in our experiments. We also discuss quantization-based compression methods in a general formulation.Notations. We use a non-bold letter like a or A to denote a scalar, a boldfaced lower-case letter like a to denote a vector, and a boldfaced upper-case letter such as A to denote a matrix. u = [u i ] T means that the i-th element of column vector u is u i and U = [u i ] means the i-th column vector of matrix U is u i . Let A be a positive definite (PD) matrix and s ∈ R, we define A s = U Λ s U T , where U ΛU T is the Singular Value Decomposition (SVD) of A. tr(A) represents the trace of a matrix A. The inner product of two matrices A and B is denoted as A, B = tr(A T B). The Frobenius norm of a matrix A is A F = A, A . A B means the elementwise matrix product (Hadamard product).Diag(a) is a diagonal matrix with diagonal vector a, while diag(A) means the diagonal vector of matrix A.The update rule of Shampoo in the matrix case combined with a first-order optimizer F isShampoo(W t-1 , L t-1 , R t-1 , s t-1 , G t ) =              L t = L t-1 +G t G T t R t = R t-1 +G T t G t G t = L -1/4 t G t R -1/4 t G t = G t ( G t F / G t F ) W t , s t = F(W t-1 , s t-1 , G t ) (1)where W t is the model parameters in matrix form, L t and R t are called preconditioners, s t is the optimizer state of F, and G t is the gradient at W t-1 . Note that L t , R t , L -1/4 t , and R -1/4 t are PD matrices. The penultimate step in (1) is the grafting trick [1] , which enables Shampoo to roughly apply the well-tuned learning rate schedule of F. The optimization variable W t does not represent all model parameters. It denotes a tensor of the model [18] or one block of a tensor [2] . In practice, we adopt an efficient and effective implementation of Shampoo for training DNNs following [2, 41] as described in Algorithm 4. In order to achieve efficient training, L t , R t , L -1/4 t , and R -1/4 t are computed once every few hundred iterations. In this case, besides L t and R t , their inverse 4-th roots should also be stored in memory, as computing them is computationally expensive. So training large models with Shampoo can be memory-intensive, consuming a significant amount of memory.Quantizing updated optimizer states using a quantizer and then dequantizing them with a dequantizer prior to use is an effective method for conserving memory. We focus exclusively on vectors, as tensors can be reshaped into vectors.Quantization. According to the idea in [8, 26] , a b-bit quantizer Q for p-dimensional real vectors is a mapping given by Q = (I • N , M) : R p → T p b × R p , where N is a normalization operator on R p , I is an elementwise function mapping any real number to an element of T b = {0, 1, . . . , 2 b -1}, and M is a maximum operator on R p . For any x ∈ R p , N and M satisfy N (x) M(x) = x.A normalization operator N for p-dimensional vectors is a transformation on R p . It scales each element of a vector x ∈ R p into [-1, 1]. A block-wise normalization operator for a p-dimensional vector x = [x 1 , x 2 , . . . , x p ] T is defined asN (x) i = x i max j∈Xi {x j } ,where N (x) i is the i-th element of N (x), and X i is a set satisfying i ∈ X i ⊂ {1, . . . , p}. Usually, X i should also satisfy X i = X j or X i ∩ X j = ∅ for i, j ∈ {1, . . . , p}. In this case, for any x ∈ R p , the number of different elements in M(x) is equal to the number of elements in set {X i |i = 1, . . . , p}.Meanwhile, the number of the elements in X i for any i should be as close as possible to a value called block size.The mappingI for x ∈ R in a b-bit quantizer Q is defined as I(x) = argmin j∈T b |x -R(j)| ,where R named quantization mapping is an elementwise function that maps any element in T b into [-1, 1], and | • | is the absolute operator for a scalar. There are three typical quantization mappings: linear quantization, dynamic quantization, and quantile quantization. Their specifications and visualizations can be found in [8] .Dequantization. Given a b-bit quantizer Q = (I • N , M) for a p-dimensional real vector x ∈ R p , the corresponding dequantizer D is a mapping defined asD(Q(x)) = D(I • N (x), M(x)) = R(I • N (x)) M(x) : T p b × R p → R p . 
methodology: 
 In this section, we describe the design of our quantization-based compression method to realize 4-bit Shampoo with fast and high precision quantization. Let Q = (I • N , M) be a quantizer and D be its corresponding dequantizer as described in Subsection 2.2.A naive approach to realize 4-bit Shampoo is applying the compression methods proposed in [8, 26] toL t , R t , L -1/4 t, and R -1/4 t in Shampoo (see (1) ). A slightly improved approach is to quantize the four PD matrices excluding their diagonal elements, which are typically much larger than their nondiagonal counterparts due to the non-negativity of the elements in diag(G t G T t ) and diag(G T t G t ). However, the naive approach can cause large quantization errors at 4-bit precision. This is because the quantization errors (or called perturbations) of quantizing L t and R t will transfer to L -1/4 t and R -1/4 t . To verify this, we first introduce two criteria to evaluate the quantization errors of matrices. We do not use the elementwise criterion in [8] . Let A denote a 32-bit matrix, g represent a transformation (can formed by quantization), and f stand for a mapping, e.g., f (A) = A -1/4 . Then we define the normwise relative error (NRE) and angle error (AE) in f of g at A asNRE = f (A) -f (g(A)) F f (A) F , AE = arccos f (A), f (g(A)) ( f (A) F f (g(A)) F .We choose two PD matrices of order 1200. The first one A 1 is derived from the real world. It is a preconditioner in 32-bit Shampoo combined with AdamW for training a Swin-Tiny model. The second one A 2 = U ΛU T is synthetic, constructed from a random orthogonal matrix U and a diagonal matrix Λ with only two distinct diagonal values. Table 1 shows the quantization errors in f (A) = A -1/4 of the naive approach at these two matrices, which are remarkably high. More analyses are given in Appendix D. The key point is that the singular values of A i (i = 1, 2) follow a specific distribution (see Figure 2 ). In this scenario, a slight perturbation of A i will significantly alter its small singular values, resulting in a drastic change to A -1/4 i .To address this issue, we propose quantizing the eigenvector matrix of a preconditioner in Shampoo, rather than the preconditioner itself. Namely, a preconditioner A is a PD matrix, and its SVD is U ΛU T , where U represents the eigenvector matrix and Λ denotes the singular value matrix. Given that Λ is a diagonal matrix, we can focus on quantizing U using Q while leaving Λ unchanged.From Table 1 , one can observe that quantizing U can significantly reduce the quantization errors. We will theoretically discuss the advantages of quantizing U compared to quantizing A in Section 4.In practice, the randomized SVD method [19] is adopted to compute the SVD of A efficiently, as shown in [40] . We want to highlight that quantizing the original L t and R t in Shampoo involves significant computational burdens to compute their inverse 4-th roots L -1/4 t and R-1/4 t , whereas quantizing the eigenvector matrices of L t and R t allows for rapid inverse root calculation. So the computational time required for both approaches is comparable (see Figure 1 ). Table 1 : Quantization errors in A -1/4 of different quantization schemes at a PD matrix A. We employ block-wise normalization with a block size of 64. U is the eigenvector matrix of A, QM = quantized matrix, and OR = orthogonal rectification. A = A 1 Synthetic A = A 2 Mapping R Bit QM OR NRE ↓ AE ( • ) ↓ Mapping R Bit QM OR NRE ↓ AE ( • ) ↓ DT 8 A 0.(V t2 Λ s V T t2 ) -1/s (V t2 ΛV T t2) and identity matrix I. Mean errors are shown on a log 10 scale.Let A be a PD matrix with SVD U ΛU T . Note that the eigenvector matrix U is orthogonal, whereas V = D(Q(U )) may not be. To further mitigate the quantization errors mentioned in Subsection 3.1, we propose employing Björck orthonormalization [4] to orthogonalize V . Particularly, given V 0 = V , we iterateEQUATIONfor t 1 ≥ 1 times and take V t1 as the rectified result. Equation (2) can also be interpreted as the gradient descent of problem min V V T V -I 2 F using a step size of 0.5, where I denotes the identity matrix. We empirically find that only one iteration (i.e., t 1 = 1) is enough. Table 1 illustrates the benefit of rectifyingV into V 1 .The update frequencies for the preconditioners and their inverse 4-th roots differ (see Algorithm 3). Given V and Λ, we also require orthogonal rectification to compute A s rapidly for any s ∈ R. The reason is as follows. It is easy to compute A s = U Λ s U T by definition. However, U Λ s U T can be very sensitive to the orthogonality of U for s < 0, making V Λ s V T largely deviate from (V ΛV T ) s ≈ A s . Similarly, we can approximate A s by V t2 Λ s V T t2 , where V t2 is generated by (2) . Figure 3 illustrates the elementwise mean errors between (V t2 Λ s V T t2 ) -1/s (V t2 ΛV T t2 ) and I for various s and t 2 , where A is the real-world matrix used in Table 1 . Based on the observation from Figure 3 , we set t 2 = 4 in our experiments.The quantizer Q is defined by the normalization operator N and mapping R, and N is determined by X i . Since an eigenvector has a unit length, the elements in X i should belong to the same column of an eigenvector matrix, i.e., they are from the same eigenvector. Instead of employing dynamic tree (DT) quantization as mapping R, we recommend utilizing linear square (Linear-2) quantization as R, particularly when b = 4. Linear-2 quantization is defined asEQUATIONwhere j ∈ T b = {0, 1, . . . , 2 b -1}. As shown in Table 1 , Linear-2 quantization has lower quantization errors compared to DT quantization at 4-bit precision.We first describe the update processes of the preconditioners and their inverse 4-th roots in our 4-bit Shampoo. A preconditioner A is a PD matrix and its SVD is U ΛU T . We can compress A into a pair (λ, U ) = (diag(Λ), Q(U )) and decompress it into (Λ, V ) = (Diag(λ), D(U )). Λ = Diag(λ), V = D(U ) 2: Rectify V by iterating (2) t 1 times 3: A = βV ΛV T + (1-β)M 4: Compute A = P ΣP T by randomized SVD 5: return diag(Σ), Q(P )Algorithm 2 PIRU(λ, U ) Input: singular value vector λ, quantized eigenvector matrix U , number of iterations t 2 for rectification, dampening term I, Q and D1: Λ = Diag(λ), V = D(U ) 2: Rectify V by iterating (2) t 2 times 3: A = V (Λ + max{λ} I) -1/4 V T 4: a = diag( A) 5: return a, Q( A -Diag(a))Algorithm 3 Practical 4-bit Shampoo Input: W 0 ∈ R m×n , L 0 = I m , R 0 = I n , L 0 = I m , R 0 = I n , β ∈ (0,λ 0,L = diag(L 0 ), U 0,L = Q(I m ); λ 0,R = diag(R 0 ), U 0,R = Q(I n ) 2: l 0 = diag( L 0 ), L 0 = Q(0); r 0 = diag( R 0 ), R 0 = Q(0) 3: for t = 1, 2, . . . , T do 4:Receive loss function L t : R m×n → R and compute gradient G t = ∇L t (W t )if t%T 1 ≡ 0 then 6:λ t,L , U t,L = PU(λ t-1,L , U t-1,L , G t G T t ); λ t,R , U t,R = PU(λ t-1,R , U t-1,R , G T t G t ) 7:else 8:λ t,L , U t,L = λ t-1,L , U t-1,L ; λ t,R , U t,R = λ t-1,R , U t-1,R 9:if t%T 2 ≡ 0 then l t , L t = l t-1 , L t-1 ; r t , R t = r t-1 , R t-113:L t = Diag(l t ) + D(L t ); R t = Diag(r t ) + D(R t ) 14: G t = L t G t R t ; G t = G t ( G t F / G t F ) 15: W t , s t = F(W t-1 , s t-1 , G t ) 
theoretical analysis: 
 In this section, we analyze why quantizing the eigenvector matrix of a preconditioner in Shampoo is better than quantizing the preconditioner itself under a certain singular value distribution. Furthermore, we consider quantization as a perturbation and prove the convergence of the perturbed Shampoo (Algorithm 6) in Appendix E. The following lemma reveals some good properties of perturbing the eigenvector matrix of a PD matrix. Lemma 1. Let A be a PD matrix whose SVD is U ΛU T , where U = [u i ] is an orthogonal matrix and Λ = diag([λ i ] T ) is a diagonal matrix. Given a perturbation ∆U = [∆u i ] and s ∈ R, we define B := (U ΛU T ) s and ∆B := ((U +∆U )Λ(U +∆U ) T ) s -B.(1) If U +∆U is orthogonal and there exists α ∈ R such that ∆u i 2 ≤ α, then ∆B F B F ≤ 2α.(2) If U +∆U is orthogonal and there exists β ∈ R such that u i ,u i +∆u i ≥ 1-β ≥ 0, then B, B +∆B B F B +∆B F ≥ (1-β) 2 .From Lemma 1, it is evident that the normwise relative error and angle error in f (A) = A s of perturbing U at A = U ΛU T are independent of Λ and s. Moreover, these errors are well-bounded under some mild conditions. Empirically, for 4-bit quantization, α = 0.1 and β = 0.005 roughly meet the conditions of Lemma 1, leading to ∆B F B F ≤ 0.2 and B,B+∆B B F B+∆B F ≥ 0.99. It is very complicated to generally analyze the perturbation in f (A) = A s of perturbing A. Thus, we focus on perturbing the singular values of A. For simplicity, we assume that both A and A + ∆A have only two distinct singular values, where ∆A is a perturbation of A. The following lemma gives the perturbation in A s of perturbing the smaller singular value of A. Lemma 2. Let A be a PD matrix of order m+n whose SVD is U ΛU T , where m, n ∈ N + , n = lm, U = [u i ] is an orthogonal matrix andΛ = diag([λ i ] T ) is a diagonal matrix. Assume that Λ = diag([cλ1 T m×1 , λ1 T n×1 ] T ), c ≥ 1, and λ > 0. Given a perturbation ∆Λ = diag([0 T m×1 , ∆λ T n×1 ] T) and s ∈ R, we define B := (U ΛU T ) s and ∆B := (U (Λ+∆Λ)U T ) s -B.(1) If ∆λ n×1 = (k -1)λ1 n×1 where k > 0, then∆B F B F = √ l|k s -1| √ c 2s + l = h 1 (s, l).Moreover, h 1 (s, l) decreases monotonically with s over (-∞, 0) and increases monotonically with l over (0, +∞). (2) If ∆λ n×1 = (tc -1)λ1 n×1 where t > 0, thenB, B + ∆B B F B + ∆B F = lt s + c s (1 + lt 2s )(l + c 2s ) = h 2 (l).Moreover, h 2 (l) decreases monotonically with l over (0, (c/t) s ] and increases monotonically with l over ((c/t) s , +∞). (3) If ∆λ n×1 = (tc -1)λ1 n×1 where k = tc > 0 and l = (c/t) s , then∆B F B F = |k s -1| √ k s + 1 , B, B + ∆B B F B + ∆B F = 2 2 + k s + 1/k s .Let us make some comments on the above lemma. First, from Lemma 2(1) we haveh 1 (1, l) = ∆A F A F = √ l|k-1| √ c 2 +l . If k ≥ 1, ∆A F A F = ∆Λ F Λ F is bounded by k c √ l = t √ l. Second, if k = tc ≥ 1and s < 0, one can deduce h 2 (l) ≥ lt 2s /(1 + lt 2s ) from Lemma 2(2), which indicates that a small lt 2s is needed to achieve small h 2 (l). We can set t = 0.02 to simulate 4-bit quantization. Based on Lemma 1 and Lemma 2(3), we have the following proposition. Proposition 1. Let A be a PD matrix of order m+n whose SVD is U ΛU T , where m, n∈ N + , n = lm, U = [u i ] is an orthogonal matrix, Λ = diag([cλ1 T m×1 , λ1 T n×1 ] T ), c ≥ 1000, and λ > 0. Given ∆U = [∆u i ], ∆Λ = diag([0 T m×1 , ∆λ T n×1 ] T ), and s ≤ -0.25, we define B := (U ΛU T ) s , B 1 := ((U + ∆U )Λ(U + ∆U ) T ) s , andB 2 := (U (Λ + ∆Λ)U T ) s . If U + ∆U is orthogonal, ∆u i 2 ≤ 0.1, u i , ∆u i ≥ -0.005, ∆λ n×1 = (0.02c-1)λ1 n×1 , and l = (c/0.02) s , then 2 B 1 -B F B F ≤ 0.4 ≤ B 2 -B F B F , 6 1 - B, B 1 B F B 1 F ≤ 0.06 ≤ 1 - B, B 2 B F B 2 F .Proposition 1 requires very strong assumptions. Nevertheless, it provides insight into why quantizing A can result in a greater normwise relative error and angle error in A s , compared to quantizing U . Complete proofs of Lemma 1, Lemma 2, and Proposition 1 can be found in Appendix F.Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)? Answer: [Yes] Justification: We present the implementation details of all the experiments in the Appendix. Guidelines:• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. , with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. 
experiments: 
 In this section, we compare our 4-bit Shampoo combined with SGDM or AdamW to their 32-bit counterparts, as well as the first-order optimizers on various image classification tasks. See more experimental results on image classification and natural language modeling tasks in Appendix H.Models, datasets, and hyperparameters. We train VGG19 [36] , ResNet34 [20] , ViT-Small [10] , and Swin-Tiny [28] on the CIFAR-100 [23] and Tiny-ImageNet [24] datasets with one RTX3060Ti GPU, and train ResNet50 and ViT-Base/32 on the ImageNet-1k dataset [34] with one A800 GPU. For hyperparameter settings, we mainly follow [41] to train CNNs and [25, 44] to train vision transformers. For all the tasks, we keep the common hyperparameters of optimizers the same values. See Appendix G for experimental details.Main results. We show the performance, wall-clock time, and memory cost in Table 2 . Firstorder optimizers run 1.2x to 1.5x epochs, resulting in longer wall-clock time, yet yielding lower test accuracies compared to second-order optimizers. In comparison to 32-bit Shampoo, our 4-bit Shampoo shows comparable test accuracies with differences ranging from -0.7% to 0.5%, increases in wall-clock time varying from -0.2% to 9.5%, and memory savings of 4.5% to 41%. Compared to the first-order optimizers, the memory costs of our 4-bit Shampoo only rise by 0.8% to 12.7%. This represents a significant advancement in the utilization of second-order optimizers. Following [26] , we report the total peak GPU memory consumption rather than the optimizer's peak GPU memory consumption. Our main focus is on quantizing the states for constructing preconditioners and their inverse roots, which are approximately 7x smaller for 4-bit Shampoo compared to 32-bit Shampoo (see Appendix G). Figure 4 shows the test accuracy curves on the CIFAR-100 and ImageNet-1k Ablations. We investigate the effectiveness of our proposed quantization techniques. Table 3 indicates that quantizing the eigenvector matrix of a preconditioner is crucial for b-bit (b = 3, 4) Shampoo to maintain 32-bit performance, and orthogonal rectification is highly beneficial for 3-bit Shampoo. As for quantization mapping, linear square (Linear-2) quantization is comparable to dynamic tree (DT) quantization. We further apply our 4-bit quantization techniques to K-FAC [30] , AdaBK [41] and CASPR [13] and the results are shown in Table 4 . We can see that the 4-bit optimizers match the performance of their 32-bit counterparts, and reduce memory by over 20%. Second-order optimizers. Different second-order optimizers apply different second-order information.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material? Answer: [Yes] Justification: We use publicly available datasets and will release our source code. Guidelines:• The answer NA means that paper does not include experiments requiring code.• Please see the NeurIPS code and data submission guidelines (https://nips.cc/public/ guides/CodeSubmissionPolicy) for more details. • While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark). Justification: We discuss both potential positive societal impacts and negative societal impacts at the end of the paper. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
related work: 
 Hessian-based optimizers [39, 27] use the Hessian matrix or its approximation. Fisher-based optimizers [30, 41] utilize the covariance matrix of the accumulated gradients or its approximation based on Kronecker product. Shampoo [18] and CASPR [13] approximate the full AdaGrad [12] preconditioner by a set of small preconditioning matrices.Memory efficient optimizers based on factorization. Adafactor [35] employs the outer product of two vectors to approximate the second moment of Adam [22] . SM3 [3] considers approximating the second moment of Adam by its covers' statistics. [14] and [40] reduce memory cost of the preconditioner in a second-order optimizer with its low-rank approximation through truncated SVD.Memory efficient optimizers based on quantization. Dettmers et al. [8] introduce block-wise dynamic quantization that enables the use of first-order optimizers with 8-bit states. Li et al. [26] push the optimizer states of Adam/AdamW to 4-bit. 
conclusions, limitations, and broader impact: 
 We propose 4-bit Shampoo, the first low-bit second-order optimizer, designed for memory-efficient training of DNNs. We find that quantizing the eigenvector matrix of the preconditioner is essential to minimize quantization errors in its inverse 4-th root at 4-bit precision, given its sensitivity to alterations in small singular values. We further introduce orthogonal rectification and linear square quantization mapping to improve performance. 4-bit Shampoo achieves lossless performance to 32-bit counterpart in training different DNNs on various tasks.Limitations. Preconditioners in Shampoo are symmetric matrices and can be stored as upper triangular matrices, saving almost half of the memory usage. However, the eigenvector matrix of a preconditioner is not symmetric, causing an 8-bit preconditioner to occupy the same memory as its 4-bit eigenvector matrix. Notably, a comparison of  
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: We present 4-bit Shampoo for memory efficient training of deep models. It poses no such risks. Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: We properly mention all the existing assets. Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL.• The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets? Answer: [Yes] Justification: We provide anonymized zip file of our code. Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: The implementation of 32-bit Shampoo used in our experiments is described in Algorithm 4. Our Pytorch implementation of Shampoo is partially based on the code provided by [2] . We implement CASPR by replacing 12 of Algorithm 4 and line 14 of Algorithm 3. We summarize the implementation of 32-bit K-FAC/AdaBK in Algorithm 5, where X t is the input feature and Y t is the output feature gradient. Both power iteration [5] and Schur-Newton iteration [17] are run for 10 iterations. Our implementation of 4-bit K-FAC/AdaBK is similar to 4-bit Shampoo (i.e., compressing L t , R t , L t , and R t ).G t = L t G t R t with J t = L t G t + G t R t ; G t = L t J t + J t R t in lineAlgorithm 4 Practical 32-bit Shampoo Input: initial parameter W 0 ∈ R m×n , left preconditioner L 0 = I m , right preconditioner R 0 = I n , inverse root of left preconditioner L 0 = I m , inverse root of right preconditioner R 0 = I n , total number of steps T , interval of updating preconditioners T 1 , interval of updating inverse roots of preconditioners T 2 , exponential decay rate for preconditioners β ∈ (0, 1), first-order optimizer F, first-order optimizer state s 0 = 0. Output: final parameter W T .1: for t = 1, 2, . . . , T do 2:Receive loss function L t : R m×n → R and compute gradient G t = ∇L t (W t )3:if t%T 1 ≡ 0 then 4:L t = βL t-1 + (1 -β)G t G T t ; R t = βR t-1 + (1 -β)G T t G t 5:else 6:L t = L t-1 ; R t = R t-1 7:if t%T 2 ≡ 0 then 8:Compute maximum eigenvalues λ L max and λ R max of L t and R t by power iteration 9: Compute L t = (L t + λ L max I m ) -1/4 and R t = (R t + λ R max I n ) -1/L t = L t-1 ; R t = R t-1 12: G t = L t G t R t ; G t = G t ( G t F / G t F ) 13: W t , s t = F(W t-1 , s t-1 , G t ) Algorithm 5 Practical 32-bit K-FAC/AdaBK Input: initial parameter W 0 ∈ R m×n , left preconditioner L 0 = 0, right preconditioner R 0 = 0,inverse root of left preconditioner L 0 = I m , inverse root of right preconditioner R 0 = I n , total number of steps T , interval of updating preconditioners T 1 , interval of updating inverse roots of preconditioners T 2 , , exponential decay rate for preconditioners β ∈ (0, 1), α = 1 for K-FAC / α = 2 for AdaBK, first-order optimizer F, first-order optimizer state s 0 = 0. Output: final parameter W T .1: for t = 1, 2, . . . , T do 2:Receive loss function L t : R m×n → R and compute gradientG t = ∇L t (W t ) 3:Receive X t by forward propagation and Y t by backward propagation 4:if t%T 1 ≡ 0 then 5:L t = βL t-1 + (1 -β)Y t Y T t ; R t = βR t-1 + (1 -β)X t X T t 6:else 7:L t = L t-1 ; R t = R t-1 8:if t%T 2 ≡ 0 then 9:Compute maximum eigenvalues λ L max and λ R max of L t and R t by power iteration 10: Compute L t = (L t + λ L max I m ) -1/α and R t = (R t + λ R max I n ) -1/L t = L t-1 ; R t = R t-1 13: G t = L t G t R t ; G t = G t ( G t F / G t F ) 14: W t , s t = F(W t-1 , s t-1 , G t ) 0 1 2Given an initial matrix P 0 ∈ R n×n , randomized SVD method computes the eigenvector matrix of a PD matrix A ∈ R n×n by iteratingEQUATIONwhere QR(X) denotes the QR decomposition of matrix X, returning an orthogonal matrix. Since we can initialize P 0 with the previous result (e.g., V in Algorithm 1), only a few iterations are enough to obtain an accurate estimation in practice. In our experiments, we iterate (4) once for Shampoo/CASPR, and iterate (4) twice for K-FAC/AdaBK.We present the constructions of different quantization mappings in b-bit quantizers (R in Q). See q k = (p k + p k+1 )/2,where k ∈ {0, . . . , 2 F -1}; c) p j = 0.9j/2 F + 0.1, where j ∈ {0, . . . , 2 F }. For 4-bit quantization, DT quantization maps T 4 onto {-0.8875, -0.6625, -0.4375, -0.2125, -0.0775, -0.0325, -0.0055, 0.0000, 0.0055, 0.0325, 0.0775, 0.2125, 0.4375, 0.6625, 0.8875, 1.0000}. For 3-bit quantization, DT quantization maps T 3 onto {-0.7750, -0.3250, -0.0550, 0.0000, 0.0550, 0.3250, 0.7750, 1.0000}.For 4-bit quantization, linear square (Linear-2) quantization maps T 4 onto {-1.0000, -0.7511, -0.5378, -0.3600, -0.2178, -0.1111, -0.0400, 0.0000, 0.0044, 0.0400, 0.1111, 0.2178, 0.3600, 0.5378, 0.7511, 1.0000}. For 3-bit quantization, Linear-2 quantization maps T 3 onto {-1.0000, -0.5102, -0.1837, 0.0000, 0.0204, 0.1837, 0.5102, 1.0000}.We present more quantization error analyses of the preconditioners. Recall that we define two kinds of quantization errors in mapping f of transformation g at A ∈ R m×n (in short errors in f (A) of g) in Subsection 3.1. Here we extend them as follows: define the normwise relative error (NRE) in f of (g 1 , g 2 ) at A asNRE = f (A) -g 2 • f • g 1 (A) F f (A) F ,and the angle error (AE) in f of (g 1 , g 2 ) at A asAE = arccos f (A), g 2 •f •g 1 (A) f (A) F g 2 •f •g 1 (A) F .Table 5 is an extension of Table 1 for Bit=4. Since the diagonal elements of A -1/4 are usually much larger than its non-diagonal elements where A is a PD matrix, we further consider the quantization errors in f (A) = A -1/4 -Diag(diag(A -1/4 )) at 4-bit precision as shown in Table 6 . Table 7 shows the quantization errors at 8-bit precision.A large condition number of a PD matrix A is indispensable for the superiority of quantizing U over quantizing A, where U is the eigenvector matrix of A. We consider contracting the singular value distribution of A = A 1 with SVD U Diag(λ)U T used in Table 5 by mapping each singular value λ of A to h(λ) = τ (λ -λ A min ) + λ A min , where λ A min is the minimum singular value of A and τ > 0 is the contraction coefficient. Figure 6 shows 4-bit quantization errors inA -1/4 or A -1/4 -Diag(diag(A -1/4 )) of quantizing U or A at A = U Diag(h(λ))U T .Table 5 : Quantization errors in f (A) = A -1/4 of different 4-bit quantization schemes at a PD matrix A. We employ block-wise normalization with a block size of 64. U is the eigenvector matrix of A and B = (g 1 (A)) -1/4 . QM = quantized matrices and OR = orthogonal rectification. Real-world A = A1 Synthetic A = A2 Mapping R QM OR NRE ↓ AE ( • ) ↓ Mapping R QM OR NRE ↓ AE ( • ) ↓ DT A 0.(A) = A -1/4 -Diag(diag(A -1/4)) of different 4-bit quantization schemes at a PD matrix A. We employ block-wise normalization with a block size of 64. U is the eigenvector matrix of A and B = (g 1 (A)) -1/4 . QM = quantized matrices and OR = orthogonal rectification.Real-world Table 7 : Quantization errors in f (A) of different 8-bit quantization schemes at a PD matrix A, where A = A 1 is derived from the real world as described in Subsection 3.1. We employ block-wise normalization with a block size of 256. U is the eigenvector matrix of A and B = (g 1 (A)) -1/4 . QM = quantized matrices and OR = orthogonal rectification. A = A1 Synthetic A = A2 Mapping R QM OR NRE ↓ AE ( • ) ↓ Mapping R QM OR NRE ↓ AE ( • ) ↓ DT A 0.f (A) = A -1/4 f (A) = A -1/4 -Diag(diag(A -1/4 )) Mapping R QM OR NRE ↓ AE ( • ) ↓ Mapping R QM OR NRE ↓ AE ( • ) ↓ DT A 0.We define the normwise relative error (NRE) and angle error (AE) of B deviating from A asNRE = B -A F A F , AE = arccos A, B A F B F .Consider Shampoo using 4-bit preconditioners for parameter updates, but also recording 32-bit preconditioners at the same time. We extract the left preconditioners L 4 and L 32 ∈ R 1200×1200 of a specific model parameter block W ∈ R 1200×768 every 8000 steps in the Swin-Tiny training on CIFAR-100 with AdamW+Shampoo. Here L 4 is a decompressed 4-bit preconditioner, and L 32 is a 32-bit preconditioner. is computed by Algorithm 2 without quantization where = 10 -4 , t 2 = 4. We find that = 10 -6 for Algorithm 2 used in our main experiments though is effective, yet it can cause a large numerical instability in the later stage of training (see Figure 8 ). Step (×8000) Step (×8000) 12 Step (×8000) 0.8 0More notations. Given a symmetric real matrix A, A 0 means that A is positive semidefinite (PSD), and A 0 means that A is positive definite (PD). Assume that symmetric matrices A and B are symmetric, the notations A B and A B mean that A -B 0 and A -B 0 respectively. Let A be a PSD matrix and s ∈ R, we define A s = U Λ s U T , where U ΛU T is the Singular Value Decomposition (SVD) of A. The Mahalanobis norm of a vector x induced by a PD matrix A is Algorithm 6 Perturbed Shampoo in the matrix case Input: W 0 ∈ R m×n , L 0 = 0 m×m , R 0 = 0 n×n , ρ 0 = 0, µ 0 = 0.x A = √ x T Ax. The dual norm of • A is denoted by • * A , where x * A = √ x T A -1 x. The spectral norm of matrix A is A 2 = sup x =0 { Ax 2 / x 2 }. A ⊗ B1: for t = 1, . . . , T do 2:Receive loss function:f t : R m×n → R 3:Compute gradient:G t = ∇f t (W t ) 4:Update preconditioners:J t = L t-1 + G t G T t ; K t = R t-1 + G T t G t 5:Perturb preconditioners:L t = g(J t ); R t = g(K t ) 6:Accumulate errors:ρ t = ρ t-1 + J t -L t 2 ; µ t = µ t-1 + K t -R t 2 7:Update parameters:W t+1 = W t -η(( + ρ t )I m + L t ) -1/4 G t (( + µ t )I n + R t ) -1/4We consider quantization as a perturbation and present the perturbed Shampoo in Algorithm 6 for convergence analysis. The regret bound of the perturbed Shampoo can be found in Theorem 1.Complete proofs can be found in Appendix F. We first introduce some basic technical tools, and the details of them are in [18, 21] . Lemma 3. Let A, A , B, B be matrices of appropriate dimensions, and u, v be two column vectors.The following properties hold:(1)(A ⊗ B)(A ⊗ B ) = (AA ) ⊗ (BB ); (2) (A ⊗ B) T = (A T ⊗ B T ); (3) If A, B 0 and s ∈ R, then (A ⊗ B) s = (A s ⊗ B s ); (4) If A A and B B , then A ⊗ B A ⊗ B ;(5) tr(AB) = tr(A)tr(B);(6) vec(uv T ) = u ⊗ v. Lemma 4. Let G ∈ R m×n , L ∈ R m×m , R ∈ R n×n , then it holds that (L ⊗ R T )vec(G) = vec(LGR). Lemma 5. Assume that 0 X i Y i for i = 1, . . . , n. Assume further that all X i commute with each other and all Y i commute with each other. Let α 1 , . . . , α n ≥ 0 such thatn i=1 α i = 1, then X α1 1 • • • X αn n Y α1 1 • • • Y αn n . Lemma 6. Let 0 ≤ α ≤ 1 and 0 X Y , then X α Y α . Lemma 7. Let A 0 and B 0, then it holds that A B if and only if B -1 A -1 . Lemma 8 (von Neumann). Let A, B ∈ R m×n and q = min{m, n}. Let σ 1 (A) ≥ • • • ≥ σ q (A) and σ 1 (B) ≥ • • • ≥ σ q (B) denote the non-increasingly ordered singular values of A and B, respectively. ThenA, B ≤ q i=1 σ i (A)σ i (B).Lemma 9. Assume that function f t is continuously differentiable and convex on R d , and matrix H t 0 for t = 1, . . . , T . Given w 0 ∈ R d , η > 0, define w t+1 = w t -ηH -1 t g t , where g t = ∇f t (w t ). Then for any w * ∈ R d , we haveT t=1 f t (w t ) - T t=1 f t (w * ) ≤ 1 2η T t=1 ( w t -w * 2 Ht -w t+1 -w * 2 Ht ) + η 2 T t=1 ( g t * Ht ) 2 .Lemma 10. Let g 1 , . . . , g T be a sequence of vectors. For ρ > 0, defineH t = (ρI + t s=1 g s g T s ) 1/2 . Then we have T t=1 ( g t * Ht ) 2 ≤ 2tr( H T ).Lemma 11. Assume that G 1 , . . . , G T ∈ R m×n are matrices of rank at most r. Let s for t = 1, . . . , T . Then for any ≥ 0,I mn + 1 r T t=1 g t g T t ( I m + T t=1 G t G T t ) 1/2 ⊗ ( I n + T t=1 G T t G t ) 1/2 .The key to the convergence proof of Algorithm 6 is forming a PD matrix sequence {H i } T i=1 , which satisfies 0 ≺ H 1 • • • H T . To achieve it, we gives the following lemma extended from Lemma 2 in the Appendix of [40] . Lemma 12. Let {X t } t=T t=1 be a sequence of symmetric matrices, and A t = t s=1 X s , where t = 1, . . . , T . Suppose we have two sequences of symmetric matrices {Y t } t=T t=1 , {Z t } t=T t=0 , and a sequence real numbers {ρ t } t=T t=0 satisfyingY t = Z t-1 + X t , ρ t = ρ t-1 + Y t -Z t 2 , Z 0 = 0, ρ 0 = 0. Define B t = ρ t I + Z t ,where I denotes the identity matrix. Then for t = 1, . . . , T , we haveB t B t-1 + X t , A t B t 2ρ t I + A t .Theorem 1. Assume that the gradients G 1 , . . . , G T ∈ R m×n are matrices of rank at most r. Then for any W * ∈ R m×n and > 0, if η = D/ √ 2r, the regret of Algorithm 6 is bounded as follows,T t=1 f t (W t ) - T t=1 f t (W * ) ≤ √ 2rD[2 1/4 mρ 1/4 T + tr( L1/4 T )][2 1/4 nµ 1/4 T + tr( R1/4 T )],whereD = max t∈[T ] W t -W * F , Lt = I m + T t=1 G t G T t , and Rt = I n + T t=1 G T t G t .Though we get a convergence guarantee of Algorithm 6, the upper bound given by Theorem 1 is very slack, since 2 1/4 mρ 1/4T is about the same as tr( L1/4 T ) for 4-bit quantization schemes in practice.Lemma 1. Let A be a PD matrix whose SVD is U ΛU T , where U = [u i ] is an orthogonal matrix and Λ = diag([λ i ] T ) is a diagonal matrix. Given a perturbation ∆U = [∆u i ] and s ∈ R, we define B := (U ΛU T ) s and ∆B := ((U +∆U )Λ(U +∆U ) T ) s -B.(1) If U +∆U is orthogonal and there exists α ∈ R such that ∆u i 2 ≤ α, then∆B F B F ≤ 2α.(2) If U +∆U is orthogonal and there exists β ∈ R such that u i ,u i +∆u i ≥ 1-β ≥ 0, then B, B +∆B B F B +∆B F ≥ (1-β) 2 .Proof.(1) Since U and U + ∆U are orthogonal, we haveB = U Λ s U T , B + ∆B = (U + ∆U )Λ s (U + ∆U ) T ,by definition. This leads to∆B = U Λ s ∆U T + ∆U Λ s (U + ∆U ) T .The Frobenius norm satisfies the triangle inequality and is orthogonality invariant. Hence,∆B F = U Λ s ∆U T + ∆U Λ s (U + ∆U ) T F ≤ U Λ s ∆U T F + ∆U Λ s (U + ∆U ) T F = Λ s ∆U T F + ∆U Λ s F = 2 ∆U Λ s F = 2 i λ s i ∆u i 2 2 = 2 i λ 2s i ∆u i 2 2 ≤ 2 i λ 2s i α 2 = 2α i λ 2s i = 2α Λ s F = 2α B F .(2) Similar to (1), we have∆B = U Λ s ∆U T + ∆U Λ s U T + ∆U Λ s ∆U T . From u i , u i + ∆u i ≥ 1 -β ≥ 0, we get 0 ≥ u i , ∆u i ≥ -β ≥ -1 because 1 = u i 2 u i + ∆u i 2 ≥ u i , u i + ∆u i = 1 + u i , ∆u i ≥ 1 -β ≥ 0,holds due to the orthogonality of U and U + ∆U . Hence,B, ∆B = tr(2U Λ 2s ∆U T + U Λ s U T ∆U Λ s ∆U T ) = tr i 2λ 2s i u i ∆u T i + tr i λ s i u i u T i j λ s j ∆u j ∆u T j = i 2λ 2s i u i , ∆u i + ij λ s i λ s j u i , ∆u j 2 ≥ i 2λ 2s i u i , ∆u i + i λ 2s i u i , ∆u i 2 = i λ 2s i [(1 + u i , ∆u i ) 2 -1] ≥ i λ 2s i [(1 -β) 2 -1] = [(1 -β) 2 -1] Λ s 2 F = [(1 -β) 2 -1] B 2 F = [(1 -β) 2 -1] B, B . Therefore, we have B, B + ∆B B F B + ∆B F = B, B + ∆B B, B = 1 + B, ∆B B, B ≥ (1 -β) 2 .The proof is completed.Lemma 2. Let A be a PD matrix of order m+n whose SVD is U ΛU T , where m, n ∈ N + , n = lm, U = [u i ] is an orthogonal matrix and Λ = diag([λ i ] T ) is a diagonal matrix. Assume that Λ = diag([cλ1 T m×1 , λ1 T n×1 ] T ), c ≥ 1, and λ > 0. Given a perturbation ∆Λ = diag([0 T m×1 , ∆λ T n×1 ] T ) and s ∈ R, we define B := (U ΛU T ) s and ∆B := (U (Λ+∆Λ)U T ) s -B.(1) If ∆λ n×1 = (k -1)λ1 n×1 where k > 0, then∆B F B F = √ l|k s -1| √ c 2s + l = h 1 (s, l).Moreover, h 1 (s, l) decreases monotonically with s over (-∞, 0) and increases monotonically with l over (0, +∞). (2) If ∆λ n×1 = (tc -1)λ1 n×1 where t > 0, thenB, B + ∆B B F B + ∆B F = lt s + c s (1 + lt 2s )(l + c 2s ) = h 2 (l).Moreover, h 2 (l) decreases monotonically with l over (0, (c/t) s ] and increases monotonically with l over ((c/t) s , +∞). (3) If ∆λ n×1 = (tc -1)λ1 n×1 where k = tc > 0 and l = (c/t) s , then∆B F B F = |k s -1| √ k s + 1 , B, B + ∆B B F B + ∆B F = 2 2 + k s + 1/k s .Proof.(1) Since U is orthogonal, we have∆B F = (Λ + ∆Λ) s -Λ s F = √ n|k s -1|λ s , B F = Λ s F = mc 2s + nλ s .Hence,∆B F B F = √ n|k s -1| √ mc 2s + n = √ l|k s -1| √ c 2s + l = h 1 (s, l) ≥ 0.It is easy to check that h 1 increases monotonically with l over (0, +∞). To prove h 1 decreases monotonically with s over (-∞, 0), defineg 1 (s) = 1 l (h 1 (s, l)) 2 = (k s -1) 2 c 2s + l .Consider the derivative of g 1g 1 (s) = (c 2s + l)2(k s -1)k s ln k -(k s -1) 2 c 2s 2 ln c (c 2s + l) 2 = 2(k s -1) (c 2s + l)k s ln k -(k s -1)c 2s ln c (c 2s + l) 2 . If s < 0 and k > 1, then k s -1 < 0, k s ln k > 0 leading to g 1 (s) < 0 since c ≥ 0; Similarly, if s < 0 and 0 < k ≤ 1, then k s -1 ≥ 0, k s ln k ≤ 0 leading to g 1 (s) ≤ 0.Thus g 1 (s) is a monotonically decreasing function for s < 0, which implies that h 1 decreases monotonically with s over (-∞, 0).(2) Similar to (1), we haveB F = mc 2s + nλ s , B + ∆B F = nt 2s + mc s λ s .Besides,B, B + ∆B = tr(U Λ s (Λ + ∆Λ) s U T ) = tr(Λ s (Λ + ∆Λ) s ) = (mc 2s + nc s t s )λ 2s .Hence, we getB, B + ∆B B F B + ∆B F = nt s + mc s (m + nt 2s )(n + mc 2s ) = lt s + c s (1 + lt 2s )(l + c 2s ) = h 2 (l) ≥ 0.To prove h 2 decreases monotonically with l over (0, (c/t) s ] and increases monotonically with l over ((c/t) s , +∞), we defineg 2 (l) = (h 2 (l)) 2 = (lt s + c s ) 2 (1 + lt 2s )(l + c 2s ) ,whose monotonicity is equivalent to that of h 2 for l > 0. Consider the derivative of g 2g 2 (l) = t 2s l 2 + 2t s c s l + c 2s t 2s l 2 + l + t 2s c 2s l + c 2s = (t s -t 2s c s ) 2 l 2 -(c s -t s c 2s ) 2 (t 2s l 2 + l + t 2s c 2s l + c 2s ) 2 .If s = 0 or tc = 1, then g 2 (l) ≡ 1. If s = 0 and tc = 1, then (t s -t 2s c s ) 2 > 0, (c s -t s c 2s ) 2 > 0.In this case, let g 2 (l) = 0, we gett 2s (1 -t s c s ) 2 l 2 = c 2s (1 -t s c s ) 2 ,which implies that l = (c/t) s . It is easy to see that g 2 decreases monotonically with l over (0, (c/t) s ] and increases monotonically with l over ((c/t) s , +∞).(3) According to (1)(2), we can easily get∆B F B F = |k s -1| √ k s + 1 , B, B + ∆B B F B + ∆B F = 2 2 + k s + 1/k s .The proof is completed.Proposition 1. Let A be a PD matrix of order m+n whose SVD is U ΛU T , where m, n∈ N + , n = lm, U = [u i ] is an orthogonal matrix, Λ = diag([cλ1 T m×1 , λ1 T n×1 ] T ), c ≥ 1000, and λ > 0. Given ∆U = [∆u i ], ∆Λ = diag([0 T m×1 , ∆λ T n×1 ] T ), and s ≤ -0.25, we define B := (U ΛU T ) s , B 1 := ((U + ∆U )Λ(U + ∆U ) T ) s , andB 2 := (U (Λ + ∆Λ)U T ) s . If U + ∆U is orthogonal, ∆u i 2 ≤ 0.1, u i , ∆u i ≥ -0.005, ∆λ n×1 = (0.02c-1)λ1 n×1 , and l = (c/0.02) s , then 2 B 1 -B F B F ≤ 0.4 ≤ B 2 -B F B F , 6 1 - B, B 1 B F B 1 F ≤ 0.06 ≤ 1 - B, B 2 B F B 2 F .Proof. According to Lemma 1, we haveB 1 -B F B F ≤ 0.2, B, B 1 B F B 1 F ≥ (1 -0.005) 2 ≥ 0.99.On the other hand, from Lemma 2(3), we getB 2 -B F B F = |x -1| √ x + 1 = f 1 (x), B, B 2 B F B 2 F = 2 2 + x + 1/x = f 2 (x),where x = (0.02c) s ∈ (0, 20 -1/4 ]. It is easy to verify that f 1 decreases monotonically and f 2 increases monotonically for 0 < x < 1. Hencef 1 (x) ≥ f 1 (20 -1/4 ) ≥ 0.4, f 2 (x) ≤ f 2 (20 -1/4 ) ≤ 0.94.The proof is completed.Lemma 12. Let {X t } t=T t=1 be a sequence of symmetric matrices, and A t = t s=1 X s , where t = 1, . . . , T . Suppose we have two sequences of symmetric matrices {Y t } t=T t=1 , {Z t } t=T t=0 , and a sequence real numbers {ρ t } t=T t=0 satisfyingY t = Z t-1 + X t , ρ t = ρ t-1 + Y t -Z t 2 , Z 0 = 0, ρ 0 = 0.Define B t = ρ t I + Z t , where I denotes the identity matrix. Then for t = 1, . . . , T , we haveB t B t-1 + X t , A t B t 2ρ t I + A t .Proof. Note that for any symmetric matrix S, it holds that S 2 I S. Then we have(ρ t -ρ t-1 )I + Z t = Y t -Z t 2 I + Z t Y t .Adding ρ t-1 I on both sides, we getB t = ρ t I + Z t ρ t-1 I + Y t = ρ t-1 I + Z t-1 + X t = B t-1 + X t . Hence B t = t s=1 (B s -B s-1 ) t s=1 X s = A t .On the other hand, we haveZ t Z t -Y t 2 I + Y t = (ρ t -ρ t-1 )I + Y t .Adding ρ t I on both sides, we getB t = ρ t I + Z t (2ρ t -ρ t-1 )I + Y t = 2(ρ t -ρ t-1 )I + ρ t-1 I + Z t-1 + X t = B t-1 + 2(ρ t -ρ t-1 )I + X t .HenceB t = t s=1 (B s -B s-1 ) t s=1 2(ρ s -ρ s-1 )I + t s=1 X s = 2ρ t I + A t .The proof is completed.Theorem 1. Assume that the gradients G 1 , . . . , G T ∈ R m×n are matrices of rank at most r. Then for any W * ∈ R m×n and > 0, if η = D/ √ 2r, the regret of Algorithm 6 is bounded as follows,T t=1 f t (W t ) - T t=1 f t (W * ) ≤ √ 2rD[2 1/4 mρ 1/4 T + tr( L1/4 T )][2 1/4 nµ 1/4 T + tr( R1/4 T )],whereEQUATIONRt = I n + T t=1 G T t G t .Proof. Define Lt = ( + ρ t )I m + L t , Rt = ( + µ t )I n + R t . According to Lemma 12, Lt and Rt are positive definite. Recall the update performed in Algorithm 6,W t+1 = W t -η L-1/4 t G t R-1/4 t . For t > 0, let H t = L1/4 t ⊗ R1/4 t , g t = vec(G t) and w t = vec(W t ). Due to Lemma 3(3) and Lemma 4, we havew t+1 = w t -ηH -1 t g t . Lemma 12 implies 0 ≺ L1 • • • LT , 0 ≺ R1 • • • RT .Thus, according to Lemma 3(3)(4) and Lemma 6, we get0 ≺ H 1 • • • H T .Let H 0 = 0. By invoking Lemma 9 and Lemma 8, we obtain the regret boundT t=1 f t (W t ) - T t=1 f t (W * ) ≤ 1 2η T t=1 (w t -w * ) T (H t -H t-1 )(w t -w * ) + η 2 T t=1 ( g t * Ht ) 2 ≤ D 2 2η T t=1 tr(H t -H t-1 ) + η 2 T t=1 ( g t * Ht ) 2 = D 2 2η tr(H T ) + η 2 T t=1 ( g t * Ht ) 2 ,whereD = max t∈[T ] w t -w * 2 = max t∈[T ] W t -W * F and w * = vec(W * ).Define H t = (r I + t s=1 g s g T s ) 1/2 . Lemma 11 and Lemma 12 imply thatH t √ r L1/4 t ⊗ R1/4 t √ rH t .Using Lemma 7 and Lemma 10 along with the above equation, we obtainT t=1 ( g t * Ht ) 2 ≤ √ r T t=1 ( g t * Ht ) 2 ≤ 2 √ rtr( H T ) ≤ 2rtr(H T ).Consequently, using Lemma 3 (5) and Lemma 12, we get the desired regret boundT t=1 f t (W t ) - T t=1 f t (W * ) ≤ D 2 2η + ηr tr(H T ) = √ 2rDtr( L1/4 T )tr( R1/4 T ) ≤ √ 2rD[2 1/4 mρ 1/4 T + tr( L1/4 T )][2 1/4 nµ 1/4 T + tr( R1/4 T )], by choosing η = D/ √ 2r.The proof is completed.We use one RTX3060Ti GPU under the PyTorch 2.0.1+CUDA11.8 framework for DNN training on the CIFAR-100 and Tiny-ImageNet datasets, use one A800 GPU under the PyTorch 2.0.1+CUDA11.7 framework for DNN training on the ImageNet-1k and C4 datasets, and use two NVIDIA L40S GPUs under the PyTorch 2.0.1+CUDA11.8 framework for DNN training on the OWT dataset. To obtain the total peak memory consumption per GPU, we call "torch.cuda.max_memory_allocated".We set "torch.backends.cudnn.benchmark" to "False" for all the experiments, except when training ViT-Base/32 on the ImageNet-1k dataset. We report the total memory consumption instead of the memory consumption of the second-order optimizer. This total memory includes data, model parameters, activations, gradients, states forming the preconditioners and their inverse roots, states for the used first-order optimizer, and memory fragments. Our focus lies in quantizing the states for constructing preconditioners and their inverse roots, which are approximately 7x smaller for 4-bit Shampoo compared to 32-bit Shampoo. Because the block size is 64, its maximum value should be calculated every 64 elements and saved as a 32-bit value, resulting in an additional overhead of 0.5 bits (32/64). Consequently, the memory savings are approximately 7 times, calculated as 32/(4+0.5).In the future, we may adopt double quantization [9] to further reduce memory consumption.For SGDM, Adagrad or AdamW used in second-order optimizers, we use 32-bit optimizer states on image classification tasks and 16-bit optimizer states on natural language modeling tasks by default. For SGDM, we set the momentum to 0.9 and use an initial learning rate of 0.1. For Adagrad, we set = 10 -10 and use an initial learning rate of 0.01. For AdamW, we set β 1 = 0.9, β 2 = 0.999, and = 10 -8 and use an initial learning rate of 0.001. For quantization settings, we employ block-wise normalization with a block size of 64 and linear square quantization by default. Matrices with a size smaller than 4096 will not be quantized. For Shampoo and CASPR, we use = 10 -6 , β = 0.95 and t 1 = 1, t 2 = 4 by default. Shampoo and CASPR precondition blocks from large matrices and the maximum order of a preconditioner is 10000 for 130M LLAMA-2 and is 1200 for other models. For training loss, we use cross-entropy loss. For image classification tasks, automatic mixed precision is enabled except for training transformers on the CIFAR-100 and Tiny-ImageNet datasets.Settings on training CNNs on CIFAR-100 or Tiny-ImageNet. Minibatch size is set to 128. Weight decay is 0.0005. Data augmentation includes random crop and horizontal flip. For Shampoo, we set T 1 = 100 and T 2 = 500. In Section 5, we run SGDM for 300 epochs and SGDM+Shampoo for 200 epochs on the CIFAR-100 dataset. We run SGDM for 150 epochs and SGDM+Shampoo for 100 epochs on the Tiny-ImageNet dataset. We adopt the multi-step learning rate schedule (the learning rate is multiplied by 0.1 for every 30% epochs with a linear warmup at the first 5 epochs).Settings on training transformers on CIFAR-100 or Tiny-ImageNet. We set a patch size of 4 for ViT-small on the CIFAR-100 dataset, and a patch size of 8 for ViT-small on the Tiny-ImageNet dataset. For training Swin-Tiny on the CIFAR-100 dataset, we use a patch size of 2 and window size of 4. For training Swin-Tiny on the Tiny-ImageNet dataset, we use a patch size of 4 and window More optimizers. Table 10 shows results of training Swin-Tiny on CIFAR-100 with NadamW, Adagrad and Adagrad+Shampoo. One can see that Adagrad+4-bit Shampoo converges faster than Adagrad with ignorable extra memory overhead, and also has higher test accuracy. Besides, though NadamW [11] is slightly better than AdamW, it is still worse than our AdamW+4-bit Shampoo. M-FAC [15] is a matrix-free method computing inverse-Hessian vector products with many gradient copies. It is not memory-efficient for M-FAC to maintain m dense gradient copies (m = 1024 in its official code). Table 11 shows that both SGDM+32-bit Shampoo and SGDM+4-bit Shampoo enjoy much higher efficiency than M-FAC (m = 32) for training ResNet34 on CIFAR-100, and enjoy higher test accuracy. EVA [42] is a rank-one second-order optimizer and is memory-efficient. We train ResNet34 on CIFAR-100 with SGDM+EVA, but despite extensive hyper-parameter tuning, we fail to achieve acceleration over SGDM. Instead, we cite EVA's result of training VGG-19 on CIFAR-100 for 200 epochs (see Table 2 in [42] ). The test accuracies of SGDM+EVA and SGDM+Shampoo are 73% and 74.5%, respectively. Models, datasets, and hyperparameters. We train 124M GPT-2 [32] for 60k steps on the Open-WebText (OWT) dataset * following the nanoGPT codebase † with two NVIDIA L40S GPUs, and train 130M LLAMA-2 [37] for 20k steps and 350M LLAMA-2 for 60k steps on the C4 dataset [33] following [43] with one A800 GPU. See Appendix G for experimental details.Main results. We show the performance, wall-clock time, and memory cost in Table 12 , and the validation loss curves in Figure 10 . As with the vision tasks, our AdamW+4-bit Shampoo consistently outperformed AdamW and naive AdamW+4-bit Shampoo in terms of performance, and AdamW+32-bit Shampoo in terms of memory usage.Memory efficiency. We further check the memory usage by increasing token batch size for a language model, which is calculated as the batch size multiplied by the context length (see [43] ). To train LLAMA2-7B on the C4 dataset using a single A800 GPU (with a maximum memory of 81,920 * http://Skylion007.github.io/OpenWebTextCorpus. † https://github.com/karpathy/nanoGPT. MB), we set the context length to 256 and then determine the maximum batch size allowed by each optimizer. For Shampoo, the maximum order of a preconditioner for training LLAMA2-7B is 2048. In all experiments, gradient checkpointing is enabled. Table 13 summarizes the evaluation results. By comparison, the 32-bit Shampoo runs out of memory with a batch size of 2, while our 4-bit Shampoo supports a batch size of 64 for standard training and only encounters memory issues at a batch size of 128. These results clearly demonstrate that our 4-bit Shampoo significantly conserves memory compared to the 32-bit version. 128 OOM• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition. • Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material. • Theorems and Lemmas that the proof relies upon should be properly referenced. 


Decision: Shampoo+Adam demonstrates improved performance compared to vanilla Adam. However, its substantial memory consumption poses a challenge for large-scale training applications. This paper is pioneering in introducing 4-bit quantization to Shampoo, significantly reducing its memory requirements. Specifically, it innovatively quantizes the eigenvector matrix of the preconditioner rather than the preconditioner itself, thereby minimizing quantization errors. Additionally, the paper cleverly incorporates Bjorck orthonormalization to correct orthogonality, further enhancing performance.

All reviewers unanimously agreed on accepting this work. After personally reading the paper and the comments, I also recommend its acceptance due to its novelty and practicality, but I still have a concern that may need to be addressed in future work. Shampoo+Adam requires approximately 50% more memory than vanilla Adam, and this may limit its practical application in training large models. Moreover, while vanilla Adam’s memory usage can be significantly reduced using the  technique of ZeRO during LLM training, it appears challenging to apply ZeRO to Shampoo due to the need to involve all elements of gradients to compute the preconditioners.

Reviews: review: paper_summary: paper_summary: This work introduces a quantized Shampoo method aimed at memory-efficient network training. Shampoo, as a second-order optimizer, incurs additional memory demands due to its optimization states. By implementing 4-bit quantization specifically on Shampoo, this work effectively reduces memory usage. The primary innovation involves quantizing the eigenvector matrix of a preconditioner instead of the preconditioner itself, which significantly minimizes quantization errors. Additionally, this work incorporates Bjorck orthonormalization to further enhance performance. Extensive experiments across multiple models and datasets validate the effectiveness of the proposed methods.
strength_weakness: strength_weakness: I find this work to be both sound and interesting. One suggestion for improvement is to expand the reporting of memory costs beyond the total GPU memory consumption. Providing a breakdown of memory usage would clarify how the 4-bit Shampoo achieves memory-efficient training.
questions: questions: Please refer to the weakness.
limitations: limitations: The authors have effectively addressed several limitations, including the efficiency sacrifices due to the non-symmetric properties of eigenvector results, the evaluation being limited to image classification tasks, and the focus on only small-scale models.
paper_summary: paper_summary: This paper aims to reduce memory usage in second-order optimizers by compressing 32-bit optimizer states to 4-bit. The authors propose a method called 4-bit Shampoo, which quantizes the eigenvector matrix of the preconditioner rather than the preconditioner itself. This approach maintains the performance of 32-bit optimizers while significantly reducing memory requirements.
strength_weakness: strength_weakness: 1. **Limited Evaluation Scope**:  The experiments are limited to image classification tasks. Evaluating the method on a broader range of tasks, such as natural language processing or reinforcement learning, could provide a more comprehensive assessment of its effectiveness.
2. **Memory Savings Trade-off**: While the method reduces memory usage, the paper does not provide a detailed analysis of the computational overhead introduced by quantization and dequantization processes. Quantization can sometimes introduce additional computation that may offset memory savings.
3. **Quantization Error Handling**: The paper discusses quantization errors and proposes solutions, but it lacks a comprehensive analysis of how these errors impact the overall training dynamics, particularly over long training periods or across different types of neural networks.
4. **Orthogonality Rectification**: The use of Björck orthonormalization to rectify the orthogonality of the quantized eigenvector matrix is an interesting approach, but it may introduce additional computational complexity. The paper should evaluate the impact of this step on the overall efficiency.
5. **Comparison with More Baselines**: The comparison is mainly with 32-bit Shampoo and first-order optimizers. Including more baselines, especially other memory-efficient optimizers like Adagrad[1], M-FAC[2], and EVA[3], would strengthen the validity of the claims. 
6. **Lack of Theoretical Framework**:The 4-bit Shampoo paper lacks a solid theoretical foundation, particularly regarding convergence guarantees and the mathematical justification for its quantization approach. Eva, on the other hand, provides a theoretical interpretation from a trust-region optimization perspective.

[1] Duchi J, Hazan E, Singer Y. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 2011, 12(7).

[2] Frantar E, Kurtic E, Alistarh D. M-fac: Efficient matrix-free approximations of second-order information. Advances in Neural Information Processing Systems, 2021, 34: 14873-14886.

[3] Lin Zhang, Shaohuai Shi, Bo Li, Eva: Practical Second-order Optimization with Kronecker-vectorized Approximation, ICML 2023
questions: questions: Please refer to the section of weaknesses.
limitations: limitations: The authors discussed the limitations of this work on several aspects, such as memory cost of eigenvector matrix, limited scope of tasks, and limited model sizes.
paper_summary: paper_summary: The paper presents a way to use a second order optimizer with 4-bit quantization, to reduce memory usage. Second order optimizers such as Shampoo use additional memory to store preconditioners and other variables needed for computing the updates. This extra memory can prevent their usage in training very large models. The authors show that quantifying the Shampoo preconditioners results in poor performance, by comparing the -1/4 power of the preconditioner and its quantized version. They then show that keeping the matrix in its decomposed form A = QSQ' allows one to quantize (and orthonormalize) Q, and this results in preconditioners which are very close to the original preconditioners.

The authors show via experiments that their quantization results in an optimization trajectory that is very close to the unquantized version, while saving up to 40% memory, at the cost of a small runtime overhead.
strength_weakness: strength_weakness: The method should help in optimization of large networks, but no results are presented on large networks. 
The presentation has some minor issues, noted below.
questions: questions: 1. In table 1, as you noted in the conclusion, please add a row for A (8-bit), to compare with the 4 bit quantization of U.
2. Please add a short paragraph describing DT quantization, to make the presentation self-contained.
3. A graph from x \in [-1,1] to  Q(x) would be nice for each of the quantization schemes.
4. In the Algorithm 3, why do you need separate T_1 and T_2? Since you are computing the SVD at every T_1 steps, it might be not too expensive to also compute the preconditioners at that step.
5. "Singular values are scaled by log 10." --> "Singular values are shown on a log_{10} scale." would be a lot less confusing.
limitations: limitations: N/A
paper_summary: paper_summary: Quantization is applied to eigen matrix instead of directly to precond. The eigen matrix is then orthonormalized via  Björck orthonormalization to orthogonalize V.
strength_weakness: strength_weakness: Could benefit from extensions to LLMs:

I feel this work lacks experiments/analysis on LLMs. I only say this because the memory needed for 2nd order methods increases a lot more for language models due to their long context window. It is clear that this method would help in that instance but it is not clear to me how much it would help and if this quantization scheme has a performance boost over the naive quantization. In the vision tasks we are training for many epochs, for language we often do not see the same windows over and over again. 

Learning rate and Learning rate Schedule:

For ResNet34 CIFAR-100 the model accuracies take a big dip mid training. This is tell tail of a multi-step learning rate schedule (noted on line 514.) Such a schedule is really sub-optimal for SGD+M. This can often result in SGD+M lagging behind other optimizers, which in the paper ultimately lead to Shampoo having a shorter wall-clock time. A linear,cosine, or trapezoidal learning rate decay should be used to give proper performance benchmarks and hence timings. Similar can be said for ResNet-50 on Imagenet-1k.

For the ViT Imagenet 1-k the initial lr should likely be better tuned.
questions: questions: Would it be possible to re-run resnet experiments with a more proper learning rate decay? This would give more confidence in the timings. 
In terms of timing the SoTA right now is nAdamW so Shampoo should be compared to that for wall clock times.  

Ideally one would use the very recent schedule free optimization as a benchmark.

https://github.com/facebookresearch/schedule_free

 If it is possible to add these as a benchmarks with timings it would be more convincing. 

Furthermore, if one can add a 124M GPT2 model for example from here https://github.com/karpathy/nanoGPT, with timings and memory analysis to consider how the 4bit quantization proposed vs the naive and 32 bit effects performance that would also add a lot to the paper in my opinion.
limitations: limitations: As mentioned above the paper does not consider LLMs which can often have a large impact on memory requirement vs a standard vision models. 

The paper does not consider more SoTA optimizers like nAdamW or LARS, and optimization methods like EMA and lookahead.

Comments: title: Thank you for your reply.
Friendly Notes
Response to the rebuttal
Thank you for your prompt and positive feedback.
comment: We appreciate your thoughtful comments and enthusiasm for our work. Your insights will help us further improve our quality. Thank you again for your support.
Yes, of course! I did not expect the authors to be able to put together experiments with CASPR or PSGD. Was more of a note for future works. 

psgd.Affine has the ability to precondition sub-linearlly actually to take less memory than Adam. There is also XMat which is also quite memory efficient. Also timings for psgd are not optimized in torch. There is a jax implementation that seems to be quite a bit faster. 

psgd + your precond quantization would make a nice combination. psgd is actually quite robust to preconditioner noise. there is theory on this in the 2015 paper. 

Thank you again for the conversations.
Thank you for your valuable comments. Below, we provide a point-by-point response to address your concerns.



**1) Experiments on CASPR.**

Thanks. We implement CASPR by replacing line 12 of Algorithm 4 and line 14 of Algorithm 3 in our paper: $\hat{G}_t=\hat{L}_tG_t\hat{R}_t$ with 
$$
J_t=\hat{L}_tG_t + G_t\hat{R}_t; \hat{G}_t=\hat{L}_tJ_t + J_t\hat{R}_t
$$
according to Algorithm 1 in [1]. The table below shows the results of training Swin-Tiny on CIFAR-100 for 100 epochs with CASPR vs. Shampoo. From the following results, we can see that Shampoo performs close to CASPR, yet is more time efficient.

|         Optimizer         | Test Accuracy (%) | Wall-Clock Time (minutes) |
| :-----------------------: | :---------------: | :-----------------------: |
|    AdamW+32-bit CASPR     |       78.82       |           296.0           |
|     AdamW+4-bit CASPR     |       78.80       |           317.8           |
|   AdamW+32-bit Shampoo    |       79.34       |           260.8           |
| AdamW+4-bit Shampoo (our) |       78.63       |           273.3           |



**2) About PSGD and dampening matrices.**

Thanks. For our method, it does not need to save a whole dampening matrix in practice, since it can recover dampening matrix by only one scalar. Refer to line 174 of the “low-bit-SO-optimizers/optimizers/shampoo1.py” file in our attached code for details.

Regarding PSGD, we use its official PyTorch implementation (https://github.com/lixilinx/psgd_torch), and report its wall-clock time and memory usage for training ResNet34 on CIFAR-100 for 2 epochs using the psgd.Affine optimizer.  From the results, we can see that psgd.Affine costs more time and memory than SGDM+32 bit Shampoo.

|      Optimizer      | Wall-Clock Time (seconds) | Total GPU memory cost (MB) |
| :-----------------: | :-----------------------: | :------------------------: |
|     psgd.Affine     |            539            |           2814.3           |
|        SGDM         |            103            |           1281.6           |
| SGDM+32-bit Shampoo |            123            |           1903.3           |



**3) Per-layer adaptive quantization and orthogonal rectification.**

Thanks. This is an interesting area for exploration, and we plan to address it in future work, as we are unable to complete it within the remaining two days.



[1] CASPR: Combining Axes Preconditioners through Kronecker Approximation for Deep Learning. In ICLR, 2024.
I will retain my score, and look forward to your revised paper.
Great! I noticed some other reviewers talk about EVA, etc if anything you guys should consider CASPR, it gives a better curvature estimate and it should be trivial to extend the work here to CASPR. 

Still, both Shampoo and CASPR require dampening matrices which require more memory.  It would also be interesting to consider applying 4-bit quantization to [PSGD](https://github.com/lixilinx/psgd_torch/tree/master?tab=readme-ov-file#resources) which doesn't require dampening. 

Furthermore, since Kronecker-factorization of these curvatures basically means, per layer curvature, it would be interesting to study the effect of which layer's curvature can be naively quantized without performance drop and which require Björck orthonormalization. This may lead to some reduction in overhead of Björck orthonormalization. 

Best wishes
We appreciate the reviewer's prompt feedback and the increased score. We will include the comparisons with Adagrad, M-FAC, and EVA in the revision.
Thanks for the reply. My main concern is the insufficient comparison with Adagrad, M-FAC, and EVA. The authors have addressed it to some extent, so I would like to raise my score.
Many thanks. We are glad our response can address your concerns. We will include the results above and try our best to solve the other issues mentioned. Again, thank you for your insightful and positive comments!!
This is exactly what I wanted to see. 

I think the paper would benefit from these tables. 

"Our observations indicate that SGDScheduleFree shows rapid improvements in training and test accuracy during the early stages of training but ultimately fails to achieve the higher test accuracy." 

A accuracy curve for this would also be good for the final paper! I would also note that AdamSchedule free should be used. Kinda weirdly Adam schedule free is actually rsm prop schedule free but that's more a a naming thing. 

I would encourage the authors to consider a sharded version of the code if they already have not. 

Thank you for your time!

Best wishes
Thank you for your prompt and positive feedback. Below, we provide a point-by-point response to address your concerns.



**1) Check memory usage by increasing token batch size.**

Thanks. The token batch size for a language model is calculated as the batch size multiplied by the context length (see [1]). To train LLAMA2-7B on the C4 dataset using a single A800 GPU (with a maximum memory of 81,920 MB), we set the context length to 256 and then determined the maximum batch size allowed by each optimizer. For Shampoo, the maximum order of a preconditioner for training LLAMA2-7B is 2048. In all experiments, gradient checkpointing is enabled. The following table summarizes the evaluation results, where "OOM" stands for "out of memory."

|           Optimizer            | Batch Size | Memory Cost (MB) |
| :----------------------------: | :--------: | :--------------: |
|           8bit AdamW           |     64     |      60135       |
|           8bit AdamW           |    128     |      68689       |
|           8bit AdamW           |    256     |       OOM        |
|   8bit AdamW+32 bit Shampoo    |     2      |       OOM        |
| 8bit AdamW+4 bit Shampoo (our) |     64     |      74561       |
| 8bit AdamW+4 bit Shampoo (our) |    128     |       OOM        |

By comparison, the 32-bit Shampoo runs out of memory with a batch size of 2, while our 4-bit Shampoo supports a batch size of 64 for standard training and only encounters memory issues at a batch size of 128. These results clearly demonstrate that our 4-bit Shampoo significantly conserves memory compared to the 32-bit version.



**2) Schedule-free method for ResNet34.**

Thanks. The table below summarizes the results of ResNet34 on CIFAR-100. We run SGDM for both 200 and 300 epochs using cosine learning rate decay. We also run SGDScheduleFree for 200 and 300 epochs using the code from [4], with the default settings: learning rate (lr) of 1.0, weight decay of 0.0005, and 2000 warmup steps. Comparing the results in the following table, it is evident that SGDScheduleFree underperforms compared to SGDM when training ResNet34 on CIFAR-100.

| Epochs |    Optimizer    | Test Accuracy (%) | Wall-Clock Time (minutes) |
| :----: | :-------------: | :---------------: | :-----------------------: |
|  200   |      SGDM       |       79.67       |           116.0           |
|  300   |      SGDM       |       79.83       |           172.7           |
|  200   | SGDScheduleFree |       74.92       |           117.5           |
|  300   | SGDScheduleFree |       75.63       |           169.6           |

Indeed, we experimented with various hyperparameter settings for SGDScheduleFree. The following table summarizes the results of training ResNet34 on CIFAR-100 for 100 and 200 epochs (lr=learning rate, wd=weight decay, acc=accuracy). Our observations indicate that SGDScheduleFree shows rapid improvements in training and test accuracy during the early stages of training but ultimately fails to achieve the higher test accuracy.

|  lr  |   wd    | train acc at 100 epoch | train acc at 200 epoch | test acc at 100 epoch | test acc at 200 epoch |
| :--: | :-----: | :--------------------: | :--------------------: | :-------------------: | :-------------------: |
| 0.1  | 0.0005  |         99.92%         |         99.58%         |        1.000%         |        32.42%         |
| 1.0  | 0.0005  |         95.90%         |         96.87%         |        72.20%         |        74.92%         |
| 1.0  | 0.00005 |         99.58%         |         99.61%         |        1.060%         |        70.63%         |
| 4.0  | 0.0005  |         87.64%         |         88.91%         |        72.34%         |        73.82%         |
| 5.0  | 0.0002  |         93.71%         |         94.47%         |        70.01%         |        74.18%         |
|  10  | 0.0001  |         93.48%         |         94.19%         |        72.90%         |        73.73%         |
| 100  | 0.00002 |         84.49%         |         85.99%         |        74.50%         |        72.15%         |



[1] Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. Galore: Memory-efficient llm training by gradient low-rank projection. In ICML, 2024.

[4] The Road Less Scheduled. arXiv preprint arXiv:2405.15682, 2024.
I like it. Thank you for taking the time to do these experiments. As well as to provide clarity on imagenet.

LLM looks good, just so I get an intuitive feel for practicality could you guys do a tolerance test for when 4bit quantization stops helping. That is either analyze the LLM theoretically, or just train for a few iterations on the LLM with 4-bit shampoo vs standard shampoo while inverting the precond each iteration to see when Adam cant load into memory, when shampoo cant fit into mem and when 4-bit cant fit into mem. I suspect this will follow closely with the figures in your paper already but would be good to double check as the extended window of LLMs can compilate things.  

Can you add schedule-free to resnet34 (21M params) as well?

On a 8M densenet schedule free is hitting ~78% in 200 epochs so I assume with the extra capacity it might jump up a percent or two.