4-bit Shampoo for Memory-Efficient Network Training
Authors: Sike Wang
Pan Zhou
Jia Li
Hua Huang
Venue: NeurIPS 2024 poster (2024)
Keywords: memory efficient
second-order optimizer
Shampoo
quantization
Primary area: optimization_for_deep_networks
TL;DR: We propose the first 4-bit second-order optimizers, exemplified by 4-bit Shampoo, maintaining performance similar to that of 32-bit ones.

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.

4-bit Shampoo for Memory-Effcient Network Training
Sike Wang
Pan Zhou
Beijing Normal University si kewanggmaill bnu. edu
Singapore Management University
panzhou@smu. edw. edu. sg
Jia Li
Hua Huang
Beijing Normal University Fiali@bmi edu.
Beijing Normal University huahuang@bmu, edu.
Abstract
Second-order optimizers, maintaining matrix termed preconditioner, supe- rior 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 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 's eigenvector matrix. which also benefits the computation its inverse 4-th root. Besides, 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 its 32-bit counterpart while being more memory-effiiin
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]. significant part of such success is attributed to first-order optimizers such as stochastic gradient descent with momentum (SGDM) 31 and Adam [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 matrix known 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 become a predominant factor in memory usage. This is the primary obstacle hindering the widespread use of second-ord -rderr optimizers in the era of large models.
There two main attempts to reduce memory consumed by optimizer states. Factorization uses low- rank approximation to optimizer states. This strategy has been applied to first-order optimizers 35
"Code is available http:p///ttttttcoooo//eaaaaaaaas 'Corresponding author.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

2200
12500
12131
2250
223000
Cre
20000
Azandi
1400
Asuntt
Adarii
J0sc0
2052
Acamitie Stamppo
1200
Acamil/+32 tet Stampoo
Ada 32 SttSsamppc
10300
Adam///-22 S-amppo
Azamitii Shamcse [tates)
Acamill (tated
Ada Shamppr phumpo tnanes)
Ada Satrmhr poo
Shampoo
us 153 203 250 300
sif
1920 1530
2500
(a) Swin-Tiny CIFAR-100
(b) ViT-Base/32 on ImageNet-Ik
Figure I: 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 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 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.
Contributions: In this paper, we present the first second-order optimizers with 4-bit optimizer states by taking Shampoo [18] example, while preserving the performance achieved with 32-be bit optimizer states. While our focus on Shampoo, we believe that our approach could also be applied to other second-order optimizers (see Table Our main contributions are highlighted below.
Firstly, to maintain 32-bit performance, we propose quantizing the eigenvector matrix of a precondi- tioner 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 quantiza- tion [8] at 4-bit precision can significantly alter the small singular values, leading to drastic change in its inverse 4-th root and thus harming 4-bit Shampoo's '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 rise in the total wall-clock time compared to quantizing the preconditioner (see Figure T.
Secondly, we present two techniques for enhancing performance. As the eigenvector matrix of 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 '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-one order optimizer states. The superiority of our developed 4-bit Shampoo demonstrated in Figure !
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.
2 Preliminaries
In this section, we present Shampoo and its implementation in our experiments. We also discuss quantization-based compression methods general formulation.
Notations. We use a non-bold letter like a or A to denote a scalar, a boldfaced lower-case letter like denote vector, and boldfaced upper-case letter such A to denote matrix. u= [u,]" means that the i-th element of column vector u is u, and U = [u.] means the i-th column vector of matrix U u.,. Let A be a positive definite (PD) matrix and s E TR we define = where U U is the Singular Value Decomposition (SVD) of A. tr(A) represents the trace of matrix A. The inner product of two matrices A and B is denoted as (A, B) -tr(AAT (ATB). B). The Frobenius norm of matrix A is ||||| m= / (A, A). A C B means the elementwise matrix product (Hadamard product).

---

Da) ) is a diagonal matrix with diagonal vector while diag(A) (A) means the diagonal vector of matrix A.
2.1
Shampoo for Matrices
The update rule of Shampoo in the matrix case combined with a first-order optimizer F
L; L, 1 G.
R = -
Shampoo(W. --, Lt-1; -1, 81-1, G.)
-
G.=,, G,
(I)
G. (/=:G/[GGll//|IGG
where W; is the model parameters in matrix form, L, and R2 are called preconditioners, 8; the optimizer state of F, and G, the gradient at W.-1 Note that L. R., L, 1/4 and R, R/11/4 are PD matrices. The penultimate step in ① is the grafting trick (I]. which enables Shampoo to roughly apply the well-tuned learning rate schedule of F. The optimization variable W, does not represent all model parameters. denotes tensor of the model [88 one block of tensor 22. In practice, we adopt an efficient and effective implementation of Shampoo for training DNNs following 2
described in Algorithm 4 In order to achieve efficient training. Lt, R4, L, 1/4 and R, 1/4 computed once every few hundred iterations. In this case, besides L, and R, their inverse 4-th roots should also be stored in memory, computing them computationally expensive. So training large models with Shampoo can be memory-intensive, consuming a significant amount of memory.
2.2 inantization-basee Compression Methods
Quantizing updated optimizer states using 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 the idea in a mapping given by
8 26 b-bit quantizer Q for p-dimensional real vectors is
Q= OI R + T x R.,
where N is normalization operator on IP, I is an elementwise function mapping any real number to an element of Tr 1, and M is a maximum operator on IRP. For any x E R N and M satisfy (xx ((z))
normalization operator N for p-dimensional vectors is a transformation on IP. scales each element of a vector RP into [-1, 1]. A block-wise normalization operator for p-dimensional vector æ I2.......... is defined as
N( maxjex, (x,)²
where N(x); is the i-th element of (xx), and X; is a set satisfying i E (1....,,2 -P)- Usually, X; should also satisfy x;= xx X, X, = for i, E (1,....,P) In this case, for any E RP, the number of different elements in M(x)) ) is equal to the number of elements in set {% / = 1,..., p). Meanwhile, the number of the elements in X, for any should be as close as possible value called block size.
The mapping I for IR in b-bit quantizer Q is defined
I(x)= argmin ja R(i),,
jET,
where R named quantization mapping is an elementwise function that maps any element in To into [-1, 1], and the absolute operator for a sealar. 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 o N, ,,M) for p-dimensional real vector x E RP, the corresponding dequantizer
mapping defined D(O(=)) = D(I e (x), R(I N(x)) o M(x) T R -
3

---

Methodology
In this section, describe the design of our muantization-basee compression method to realize 4-bit Shampoo with fast and high precision quantization. Let Q= N M) be a quantizer and be its corresponding dequantizer described in Subsection 2.2
3.1 Quantizing the Eigenvector Matrices
naive approach to realize 4-bit Shampoo is applying the compression methods proposed in to Lt, R., Li -,1/ and R, in Shampoo (see (T). A slightly improved approach is to quantize the four PD matrices excluding their diagonal elements, which are typically much larger than their non- diagonal counterparts due to the non-negativity of the elements diag( G, and diag (G, G.).
However, the naive approach can cause large quantization errors at 4-bit precision. This is because the quantization errors (or called perturbations) of quantizing Lt and R. will transfer L, 114
and R 1/4 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 32-bit matrix, g represent transformation (can formed by quantization), and stand for mapping, e-g., f(A)= A-14 Then we define the normwise relative error (NRE) and angle error (AE) in at as [(g()) ([.
NRE ([(A) [[g((((())) ())il
AE=arccos
[[(A)]l
We choose PD matrices of order 1200. The first A1 is derived from the real world. It preconditioner in 32-bit Shampoo combined with Adam for training a Swin-Time model. The second one A2 U AA synthetie, constructed from random orthogonal matrix U and diagonal matrix with only two distinct diagonal values. Table [ shows the quantization errors in f(A) = A- of the naive approach at these two matrices, which are remarkably high. More analyses given in Appendix The key point is that the singular values of A, (i= 1, follow specific distribution (see Figure 2 In this scenario, a slight perturbation of A, will significantly alter its small singular values, resulting in drastic change to A, 1/4
To address this issue, we propose quantizing the eigenvector matrix of preconditioner in Shampoo, rather than the preconditioner itself. Namely, preconditioner A is PD matrix, and its SVD is U TAUT, where represents the eigenvector matrix and denotes the singular value matrix. Given that A is diagonal matrix, we can focus on quantizing U using Q while leaving A unchanged. From Table 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, and R, in Shampoo involves 1/4 significant computational burdens to compute their inverse 4-th roots L, L-1/4 and R 1/4 whereas quantizing the eigenvector matrices of L, and R; allows for rapid inverse root calculation. So the computational time required for both approaches is comparable (see Figure
Table 1: Quantization errors -1/4 of different quantization schemes at matrix A. We employ block-wise normalization with block size of 64. U is the eigenvector matrix of A, QM quantized matrix, and OR orthogonal rectification.
1/4
Real-world A A
Synthetic A A
Mapping R
Bit
QM
OR NRE
AE (?) Mapping R Bit
QM
OR NRE
AE (P)
8
A
X
0.2192
8.3014
8
A
X
0.1896
10.877
4
A
X
0.6241
17.319
4
A
X
0.4615
17.189
DT
DT
0.0709
4.0426
4
U
X
0.1224
7.0144
4
0.0455
2.5615
4
U
/
0.0878
4.9960
8
A
0.2164
7.9751
8
A
X
0.1310
7.4717
4
A
X
0.6243
17.293
4
A
X
0.4465
15.338
Linear-2
Linear-2
4
U
X
0.0543
3.1066
4
U
X
0.0942
5.3998
4
U
0.0343
1.9456
4
0.0669 U
3.8166

---

-1
Singllar Value
Singular Value
1/10
12
1/20
-8
U
-10
real
1.
600 Index
1000
400
1000
Index
Number Itertions b
(a) Real-world
(b) Synthetic
Elementwise mean enrors between (V, (V, AV,) and identity matrix I. Mean errors
Figure 3:
Figure 2: Singular value distributions of PD matrices (real) and their 4-bit compressions (quan) used in Table with R=DT, QM=A. Singular values are shown on log10 scale.
are shown on a log scale.
3.2 Rectifyyng the Orthogonality of Eigenvector Matrices
Let A be PD matrix with SVD UAUT. Note that the eigenvector matrix is orthogonal, whereas V = =D(((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 Vo - VV
we iterate
V
Vi-1
for I 2 times and take V,, as the rectified result. Equation 2 can also interpreted the gradient descent of problem miny V using a step size 0.5, where denotes the identity matrix. We empirically find that only one iteration (i.e., f = is enough. Table Tillustrate the benefit of rectifying V into V1.
The update frequencies for the preconditioners and their inverse 4-th roots differ (see Algorithm 3) Given V and A, we also require orthogonal rectification compute A rapidly for any TK. The reason is as follows. It is easy to compute A = U U by definition. However, U can be very sensitive the orthogonality of U for s making VA largely deviate from
(VAVTy ~ A Similarly, we can approximate A' by V,,'V, where Vt, is generated by (2). and for
Figure 3 illustrates the elementwise mean errors between (V, VV. (VTT)"((VV,,VV
various and f2, where A is the real-world matrix used in Table [ Based on the observation from Figure set t2 4 in experiments.
3.3 Selecting the Quantizer
The quantizer defined by the normalization operator N and mapping R, and A is determined by X; Since an eigenvector has unit length, the elements in X, should belong the same column of an eigenvector matrix, i.e., they are from the same eigenvector. Instead employing dynamic tree (DT) quantization as mapping R, we recommend utilizing linear square (Linear-2) quantization as R. particularly when b= Linear-2 quantization defined as
(-1 + 2j/(22-11) -
j 2 1;
R())
j 2-
1 (21+2/((2-1))
j> 201
where TT (0, (0,1,.... I...., -1. As shown in Table Linear-2 quantization has lower quantization errors compared to DT quantization at 4-bit precision.
3.4 Overall Algorithm
We first describe the update processes of the preconditioners and their inverse 4-th roots in our 4-bit Shampoo. A preconditioner A is PD matrix and its SVD is U T We can compress A into a pair (A, U) (diag(A) A). ((U)) and decompress it into (A, V) = (Diag(A), (U)). Algorithm (Preconditioner Update, PU) shows the update rule of A. Similarly, we compress A ~ A 1/4 into pair (a, A) (diag(A) Q(A - (a-iaag())) and decompress it into Diag(a) (A)). Algorithm
5

---

(Preconditioner's Inverse 4-th Root Update, PIRU) gives the update rule of A. Based on the above update rules, we can summarize our 4-bit Shampoo Algorithm Note that we omit some input parameters of PU and PIRU because they can be found in Algorithm 3 in the same form.
Algorithm 1 PUA M)
Algorithm 2 PIRU(A, U)
Input: singular value vector X, quantized eigen- vector matrix U, M, number of iterations I for rectification, exponential decay rate
Input:
singular value vector A, quantized eigen- vector matrix U, number of iterations for rectification, dampening term el, Q and D
B E (0, 1), 2 and D
A Diag (A),
D
1: A Diag( (2), V UD)
2: Rectify V by iterating 2 b2 times
2: Rectify V by iterating (2) t times
Å (A + max ma -1/4T
3: A AVAV (1 -1-5)M M
4: Compute A PEP by randomized SVD
diag(A) (A)
5: return diag(2) ((P)
5: return a, Q(Ā Diag(a))
Algorithm 3 Practical 4-bit Shampoo
Input: Wo E Rm Lo = elm. Rg = eIn, Lo = Im. Ro -n. 3 E (0, 1), f1, t2, update interval Ti, update interval T2, total number of steps T, first-order optimizer F, first-order optimizer state Su = 0, 4-bit quantizer Q and its corresponding dequantizer D.
Output: final parameter WT.
1: Ao diag Lo). Uo Q(Im);
70,= diag( Ro), UO Q(I,)
2: 4 diag( Lo), Lo 2(0): To diag( Ro) R. Q(0)
3: for 1 ...,, T do
4: Receive loss function L. R m + R and compute gradient G. = (W,,
5: if 1TT, = 0 then
6: AL, (LL= U-LL, G;G,); A U
U, G. Gt)
7:
else
8:
L. U.L = A U-LLL N U.R = A U.-LR
9:
if IT, = 1 then
10:
4, L. PIRU(A U,L);;
T1, R = PI U.R)
11:
else
12:
4, L, = 4. -, To. R. - Rr
13:
L. = Diag(h.) + +(I);
R Diag(()) + D(R.) R.)
14:
G, L, G. =GG(IG/
W,, 8,
Theoretical Analysis
In this section, we analyze why quantizing the eigenvector matrix of preconditioner in Shampoo is better than quantizing the preconditioner itself under certain singular value distribution. Furthermore, we consider quantization as perturbation and prove the convergence of the perturbed Shampoo (Algorithm ()) Appendix E The following lemma reveals some good properties of perturbing the eigenvector matrix of PD matrix.
Lemma Let be PD matrix whose SVD is UAT,, where U [u:] an orthogonal matrix and A =diag =diag((>.]]] ) is diagonal matrix. Given a perturbation AU [Aui] and IR, we define : (UAUT) and : (B=
- B.
(1) UUU+AU is orthogonal and there exists a R such that || 12 S a, then
| F < 2a.
|B
(2) Uf + AU is orthogonal and there exists / E IR. such that (ui, u: + Aus) 2 1- 2 0, then
(B, B BB)
B + AB 2 (1-3)2.
B
F

---

From Lemma 1 it is evident that the normwise relative enror and angle error in f(A = A' of perturbing U at A = UAUT are independent of and 8. Moreover, these errors are well-bounded under some mild conditions. Empirically, for 4-bit quantization, a 0.1 and 0.005 roughly meet the conditions of Lemma leading to ||B | S 0.2 and /B B+AB (B.B-AB) 2 0.99.
It is very complicated to generally analyze the perturbation in f(A) = A of perturbing A. Thus, we focus on perturbing the singular values of A. For simplicity, we assume that both A and A + AA have only two distinct singular values, where A is a perturbation of A. The following lemma gives the perturbation in A of perturbing the smaller singular value of A.
Lemma 2. Let A be PD matrix of order mn++n whose SVD is UAUT, where m. n. E N4, n. Im, U= = [u:] is an orthogonal matrix and A diag([Ad] )) is a diagonal matrix. Assume that = diag([ m Axxxl"") c 2 1, and X > Given a perturbation AA diag( mxx1 AXxx1] IT) and 8 I., we define B = ) and
(1) If AAxx1 = (k 1)d1ax1 where k > 0, then
|ABI |F | Vc +1
Vith h1 (s,())
Moreover, h1 (s, decreases monotonically with over (-00, 0) and increases monotonically with over (0, +00). (2) If AAnx1 (tc- 1)Alnx1 where t >> O, then
(B, BB,++D AB)
11 +
= h2(1).
B B + AB V(L+112)(( c25)
1122)(1
Moreover, (22 decreases monotonically with 1 over (0, (c/t)* and increases monotonically with over ((c/t)", ", +20).
(3) If AAxx (tc nxll where k = to 0 and I (c/t)", then
AB
(B. + AB
B F
V1++1'
B B AB F /2 + h- 1/k
Let us make some comments on the above lemma. First, from Lemma 1) we have h1 (1,
[AAl|E -!jk-1|
TAIlF ve++ If 1, ||Al|r
|AAlE is bounded by _ V11 lVl. Second, if k Lc 2 1 AA T
and 0, one can deduce hz(l) 2 1I120(1 1120) from Lemma 2) which indicates that small 2 is needed achieve small h2). We can set & 0.02 simulate a 4-bit quantization. Based on
Lemma and Lemma [33 we have the following proposition.
Proposition 1. Let be PD matrix of order m++n whose SVD is UAUT, where m, n N., lm, U [u.] is an orthogonal matrix, A = diag([c)] mx Al 1 c² 1000, and X >>00 0. Given AU [Au,), AA diag( mx1: AAT 1 ),,, and S -0.25, define B = :UAUU
B (U AU) + AU) and B, (U(A+ + AA)UT) U AU is orthogonal,
| 1l2 S 0.1, (u,, Au;) 0.005, A = (0.02c 1)A1 and 1= (c/0.02)*, then
B1 Bll
B: BIF
6(1 B B
(B, B))
B B2)
< 0..4 <
$0.06 <
|Brr
B B
Bell
Proposition requires very strong assumptions. Nevertheless, it provides insight into why quantizing A can result in greater normwise relative error and angle error in A", compared to quantizing U. Complete proofs of Lemma I Lemma]
2 and Proposition T can be found in Appendix
5 Experiments
In this section, we compare our 4-bit Shampoo combined with SGDM or AdamW to their 32-bit counterparts, well as the first-order optimizers 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 200, ViT-Small 10]
and Swin- Tiny (88 on the CIFAR-100 (3 and Tiny-Im (4 datasets with one RTX3060Ti GPU, and train ResNet50 and ViT-Base/2 the ImageNet- Ik dataset [34] with one A800 GPU.
7

---

Table 2: Performance, wall-clock time and memory cost on various image classification tasks. TA test accuracy, WCT wall-clock time, and TMC total GPU memory cost.
Dataset
Model
Optimizer
TA (%)
WCT (min)
TMC (MB)
SGDM
74.14
97.70
512.17
VGG19
SGDM 32-bit Shampoo
74.54
84.45
979.13
SGDM + 4-bit Shampoo
74.74
92.51
577.14
SGDM SGDM + 32-bit Shampoo
78.98
170.1
822.03
ResNet34
79.71
147.2
1441.8
SGDM + 4-bit Shampoo
79.17
155.8
908.40
CIFAR-100
Adam AdamW + 32-bit Shampoo
74.34
668.1
2720.0
ViT-Small
77.50
498.7
3252.0
AdamW + 4-bit Shampoo
77.22
510.8
2791.7
AdamW
76.69
318.6
1465.8
Swin-Tiny
Adam W + +22-bit Shampoo
79.34
260.8
2036.0
AdamW + 4-bit Shampoo
78.63
273.3
1543.9
SGDM
61.53
172.0
1062.3
VGG19
SGDM 32-bit Shampoo
63.39
136.5
1531.9
SGDM + 4-bit Shampoo
62.84
143.8
1127.3
SGDM SGDM 32-bit Shampoo
67.10
432.1
2304.0
ResNet34
67.90
313.0
2924.3
SGDM + 4-bit Shampoo
67.95
329.3
2390.4
ImageNet
AdamW Adam + 32-bit Shampoo
54.66
1274
2730.1
ViT-Small
57.11
953.9
3261.1
Adam + 4-bit Shampoo
57.15
970.3
2801.9
AdamW Adam + 32-bit Shampoo
58.77
701.9
1789.9
Swin-Tiny
61.74
565.3
2362.8
Adam W + 4-bit Shampoo
62.24
582.7
1868.1
SGDM SGDM 32-bit Shampoo
76.70
2134
11307
ResNet50
77.07
1910
11937
SGDM + 4-bit Shampoo
76.92
1970
11396
ImageNet-Ik
AdamW AdamW + 32-bit Shampoo
72.87
2190
10600
ViT WTT-Base/32
75.03
1774
12134
AdamW + 4-bit Shampoo
74.78
1770
10804
ResNeI34 on CIFAR-100
VIT-Small CIFAR-100
ResNet50 ImageNet-14
VTT-aase/32 on Image/Net-lk
80
Teck Accurary
SOOM
AdaniN
SGDM
Adamii
SGDN+32
AdamiN' Shampoo
SGDM+
AdamW-32 Shampoo
Shampoo
Shampon
Shampoo
Wallelock lime (imi))
Wall-llok Time (man)
Wal-dick Time (min)
Wall-clock Time [min]
Figure 4: Visualization of test accuracies on the CIFAR-100 and ImageNet- 1k datasets.
For hyperparameter settings, we mainly follow to train CNNs and 25 transformers. For all the tasks, we keep the common hyperparameters of optimizers the same values. See Appendix G for experimental details.
train vision
Main results. We show the performance, wall-clock time, and memory cost in Table First- order optimizers run 1.2x 1.5x epochs, resulting in longer wall-clock time, yet yielding lower test accuracies compared to second-ordes optimizers. In comparison to 32-bit Shampoo, our 4-bit Shampoo shows comparable test accuracies with differences ranging from -0.7% 0.5%. increases in wall-clock time varying from -0.2% to 9.5%, and memory savings of 4.5% to 41 1.. Compared to the first-order optimizers, the memory costs of our 4-bit Shampoo only rise by 0.8% 12.7%. This represents 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 quantizing the states for constructing preconditioners and their inverse roots, which are approximately 7x smaller for 4-bit Shampoo compared 32-bit Shampoo (see Appendix G Figure 4 shows the test accuracy curves the CIFAR-100 and ImageNet-

---

Table 3: Ablation study on the impact of different quantization techniques to Swin-Tiny training on the CIFAR- dataset. U the eigenvector matrix of a preconditioner A. QM quantized matrix, OR orthogonal rectification in Algorithm
TL training loss, and TA = test accuracy.
4-bit
3-bit
Mapping R
QM
OR
TI.
TA (%)
Mapping R.
QM
OR
TL.
TA (%)
Linear-2
X
1.631
76.95
Linear-2 2
A
1.648
76.70
DT
U
X
1.569
78.70
DT
U
NaN
Linear-2
U
X
1.566
78.22
Linear-2
U
X
NaN
Linear-2
U
1.551
78.63
Linear-2
U
1.572
78.53
datasets. The test accuracy curves of 4-bit Shampoo and 32-bit Shampoo are very close, both of which are above the test accuracy curves of the first-order optimizers.
Ablations. We investigate the effectiveness of our proposed quantization techniques. Table iiddicates that quantizing the cigenvector matrix of preconditioner is crucial for b-bit (b = 3, 4) Shampoo to maintain 32-bit performance, and orthogonal rectification 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 [4]] and CASPR [13] and the results are shown in Table We can see that the 4-bit optimizers match the performance of their 32-bit counterparts, and reduce memory by over 20%.
Table 4:
Performance and memory cost of
Related Work
training Swin-Tiny on CIFAR-I00 TA = test accuracy and TMC total GPU memory cost.
Second-order optimizers. Different second-order optimizers apply different second-order information. Hessian-based optimizers [39] [2 use the Hessian matrix or its approximation. Fisher-based optimiz-
Optimizer
TA (%)
TMC (MB)
AdamW+32-bi K-FAC
78.20
2388.0
Adam K-FAC
78.56
1878.3
ers [30 41 utilize the covariance matrix of the ac-
32-bit AdamW BK
79.28
2388.0
cumulated gradients or its approximation based on Kronecker product. Shampoo
4-bit AdamW BK
79.34
1878.3
[8] and CASPR [3]
Adam W+32-bit CASPR
78.82
2034.6
approximate the full AdaGrad [2 preconditioner
Adam W4i-bit CASPR
78.80
1543.9
by a of small preconditioning matrices.
Memory efficient optimizers based on factorization. Adafactor of two vectors to approximate the second moment of Adam 222] SM3 considers approximating
335
employs the outer product
the second moment of Adam by its covers' statistics. [4] 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. 22 push the optimizer states Adam/Adam to 4-bit.
7 Conclusions, Limitations, and Broader Impact
We propose 4-bit Shampoo, the first low-bit second-orde optimizer, designed for memory-efficien training of DNNs. We find that quantizing the eigenvector matrix of the preconditioner is essential to minimize quantization errors in inverse 4-th root 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 upper triangular matrices, saving almost half of the memory usage. However, the eigenvector matrix of a preconditioner is not symmetric, causing 8-bit preconditioner to occupy the same memory as its 4-bit eigenvector matrix. Notably, a comparison of Table [ and Table in Appendix [ shows that the 4-bit quantization of the eigenvector matrix has smaller quantization errors than the 8-bit quantization of the preconditioner. Our evaluation is currently limited to image classification and natural language modeling tasks. Due to limitations computing resources, we do not test our 4-bit Shampoo on large-scale models with billions of parameters.

---

Broader Impact. Our work can facilitate training large models with second-ored -oder optimizers. This could open up new research possibilities that were previously unattainable due GPU memory constraints, especially benefiting researchers with limited resources.
Acknowledgments and Disclosure of Funding
Jia Li and Hua Huang were supported by the NSF of China (grant no. 62131003). Jia Li was also supported by the NSF of China (grant no. 62102034). Pan Zhou was supported by the Singapore Ministry of Education (MOE) Academic Research Fund (AcRF) Tier grants (project ID: 23-SIS- SMU-028 and 23-SIS-SMU-070),
References
[1] Naman Agarwal, Rohan Anil, Elad Hazan, Tomer Koren, and Cyril Zhang. Disentangling adaptive gradient methods from learning rates. arXiv preprint arXiv:2002. 1803, 2020.
[2] Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer. Scalable second order optimization for deep learning. arXiv preprint arXiv:2002.09018, 2020.
[3] Rohan Anil, Vineet Gupta, Tomer Koren, and Yoram Singer. Memory efficient adaptive optimization. Advances in Neural Information Processing Systems, 32, 2019.
[4] A. Björck and C. Bowie, An iterative algorithm for computing the best estimate of orthogonal matrix. SIAM Journal on Numerical Analysis, 8(2):358-344, 1971.
[5] R.L. Burden, J.D. Faires, and A.M. Burden. Numerical Analysis. Cengage Learning, 2015.
[6] Aaron Defazio, Xingyu Yang, Harsh Mchta, Konstantin Mishchenko, Ahmed Khaled, and Ashok Cutkosky. The road less scheduled. arXiv preprint arXiv:2405. 15682, 2024.
[7] Tim Dettmers. 8-bit approximations for parallelism in deep learning. In Proceedings of the International Conference on Learning Representations, 2016.
[8] Tim Dettmers, Mike Lewis, Sam Shleifer, and Luke Zettlemoyer. 8-bit optimizers via block- wise quantization. In Proceedings of the International Conference on Learning Representations, 2022.
[9] Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, and Luke Zettlemoyer. QLoRA: Efficient finetuming of quantized LLMs. Advances Neural Information Processing Systems, 2023.
[10] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In Proceedings of the International Conference Learning Representations, 2021.
[11] Timothy Dozat. Incorporating Nesterov momentum into Adam. In Proceedings of the Interna- lional Conference on Learning Representations Workshop. 2016.
[12] John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 12(61):2121 (121):22111219 2259,, 2011.
[13] Sai Surya Duvvuri, Fnu Devvrit, Rohan Anil, Cho-Jui Hsich, and Inderjit S Dhillon. Combining axes preconditioners through Kronecker approximation for deep learning. In Proceedings of the International Conference Learning Representations, 2024.
[14] Vladimir Feinberg, Xinyi Chen, Y. Jennifer Sun, Rohan Anil, and Elad Hazan. Sketchy: Memory-cfficient adaptive regularization with frequent directions. Advances in Neural Informa- tion Processing Systems, 2023.
[15] Elias Frantar, Eldar Kurtic, and Dan Alistarh. M-FAC: Efficient matrix-free approximations of second-order information. Advances in Neural Information Processing Systems, 2021.
10

---

[16] Anmol Gulati, James Qin, Chung-Cheng Chiu, Niki Parmar, Yu Zhang, Jiahui Yu, Wei Han, Shibo Wang, Zhengdong Zhang. Yonghui Wu, and Ruoming Pang. Conformer: Convolution- augmented transformer for speech recognition. In Proceedings of the Conference of the International Speech Communication Association, 2020.
[17] Chun- Huaa Guo and Nicholas Higham. A its inverse. SIAM Journal on Matrix Analysis and Applications, 28(3):788 8088,404222 2006.
method for the matrix pth root and
[18] Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimization. In Proceedings of the International Conference on Machine Learning, 2018.
[19] N. Halko, P. G. Martinsson, and J. A. Tropp. Finding structure with randomness: Probabilistic algorithms for constructing approximate matrix decompositions. SIAM Review, 53(2):217-288, 2011.
[20] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for im- age recognition. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, June 2016.
[21] Roger Horn and Charles R. Johnson. Matrix Analysis. Cambridge university press, 2012.
[22] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Proceedings of the International Conference Learning Representations, 2015.
[23] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
[24] Ya Le and Xuan Yang, Tiny imagenet visual recognition challenge. CS 231N, 7(7):3, 2015.
[25] Seung Hoon Lee, Seunghyun Lee, and Byung Cheol Song. Vision transformer for small-size datasets. arXiv preprint 12. 2021.
[26] Bingrui Li, Jianfei Chen, and Jun Zhu. Memory efficient optimizers with 4-bit states. Advances in Neural Information Processing Systems, 2023.
[27] Hong Liu, Zhiyuan Li, David Leo Wright Hall, Percy Liang, and Tengyu Ma. Sophia: A scalable stochastic second-order optimizer for language model pre-training, In Proceedings of the International Conference on Learning Representations, 2024.
[28] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEF/CVF International Conference on Computer Vision, October 2021.
[29] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In Proceedings of the International Conference on Learning Representations, 2019.
[30] James Martens and Roger Grosse. Optimizing neural networks with Kronecker-factored approximate curvature. In Proceedings of the International Conference Machine Learning. 2015.
[31] Ning Qian. On the momentum term in gradient descent learning algorithms. Neural Networks, 151, 1999.
[32] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, and others. Language models are unsupervised multitask learners. OpenAl blog, 2019.
[33] 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):1-67, 2020.
[34] Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael S. Bernstein, Alexander C. Berg, and Li Fei- Fei. ImageNet large scale visual recognition challenge. International Journal of Computer Vision, 1 115(3):111-252, 2015.

---

[35] Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In Proceedings of the International Conference on Machine Learning. 2018.
[36] Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. In Proceedings of the International Conference on Learning Representations, 2015.
[37] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, and others. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.091 230000988.00 2023.
[38] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, E ukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in Neural Information Processing Systems, 30, 2017.
[39] Zhewei Yao, Amir Gholami, Sheng Shen, Mustafa Mustafa, Kurt Keutzer, and Michael W. Mahoney. AdaHessian: an adaptive second order optimizer for machine learning. In Proceedings of the AAAI Conference Artificial Intelligence, 2021.
[40] Jui-Nan Yen, Sai Surya Duvvuri, Inderjit S. Dhillon, and Cho-Jui Hsich. Block low-rank preconditioner with shared basis for stochastic optimization. Advances in Neural Information Processing Systems, 2023.
[41] Hongwei Yong, Ying Sun, and Lei Zhang. A general regret bound of preconditioned gradient method for DNN training. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition, June 2023.
[42] Lin Zhang, Shaohuai Shi, and Bo Li. Eva: Practical second-orde optimization with Kronecker- vectorized approximation. In Proceedings of the International Conference on Learning Repre- sentations, 2023.
[43] Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. GaLore: Memory-efiiee efficient LLM training by gradient low-rank projection. In Proceedings of the International Conference on Machine Learning, 2024.
[44] Pan Zhou, Xingyu Xie, and Shuicheng Yan. Win: Weight Weight-decay decay-integrated Nesterov accel- eration for adaptive gradient algorithms, In Proceedings of the International Conference on Learning Representations, 2023.
12

---

A Implementation Details of Shampoo, CASPR, K-FAC and AdaBK
The implementation of 32-bit Shampoo used in our experiments described in Algorithm Pytorch implementation of Shampoo is partially based on the code provided by 12. We implement CASPR by replacing G. R, with J. = L.G. G. G. L,J, + J, J, in line 12 of Algorithm 4 and line 14 of Algorithm 3 We summarize the implementation 32-bit K-FAC/AdaBK in Algorithm where X, is the input feature and Y, is the output feature gradient. Both power iteration [ and Schur-Newton iteration [17] are run for 10 iterations. Our implementation of 4-bit K-FAC/AdaBK is similar to 4-bi bit Shampoo (i.e., compressing Lt. Rs. L, and R.).
4 Our
Algorithm Practical 32-bit Shampoo
Input: initial parameter Wo |_*x", left preconditioner Lo elm, right preconditioner Ro eI inverse root of left preconditioner Lo = I inverse root right preconditioner Rg In. total number of steps T. interval of updating preconditioners Ti, interval of updating inverse roots of preconditioners T2, exponential decay rate for preconditioners / E (0, 1), first-order optimizer F, first-order optimizer state 80 0.
Output: final parameter Wr.
1: for = .,, do
2:
Receive loss function L. :RRM H R and compute gradient G; = (W,(W))
3: if t%T; = 0 then
4: L = 3L + -1- () R₂ = -B -1-
5: else
6:
L, 1,-1 R = R
7:
if 1KT2 0 then
8:
Compute maximum eigenvalues Miax and X of L, and Rg by power iteration
9:
Compute L. (L, + A (m) -1/4 m and R. (Re + A max ela) by Schur-Newton
iteration
10:
else
L - R -
11:
12:
G. = L.
G. G(GGG)))))
13:
W,, 8; St-
Algorithm 5 Practical 32-bit K-FAC/AdaBK
Input: initial parameter Wo Rmxn left preconditioner Lo 0, right preconditioner RG 0. inverse root of left preconditioner Ln Ims inverse root of right preconditioner Rs I,, total number of steps T, interval of updating preconditioners Ti, interval of updating inverse roots of preconditioners T2, e, exponential decay rate for preconditioners (0, 1). a for K-FAC = 2 for AdaBK, first-order optimizer F, first-order optimizer state 80 0.
Output: final parameter WT.
1: for 1, 2,.., do
2: Receive loss function L. : IR x + R and compute gradient G; = VL, (W,)
3: Receive X, by forward propagation and Y, by backward propagation
4: if t% = then
5:
L, =,,,, (1- R. 3R (1- (1-
6:
else
7:
L, L,1;; R
8:
if 1%72 = 0 then
9:
Compute maximum eigenvalues M and Xma of L, and R2 by power iteration
Compute L, (L, + Naax eI Im) -1/a and R2 R2 + XR iteration
max II,)) -1/0 1//x by Schur-Newton
11:
else
12:
L, - R R
G, L,
G, G
14:
8, = F(W Gi)
13

---

1.00
1.00
DT
0.75
0.75
Linear-2
Linear-2
0.50
0.50
0.25
0.25
Value
Value
0.00
0.00
-0.55
-0.25
-0.50
-0.50
-0.75
-0.75-
-1.00
-1.00
0
1
2
3
4
5
7
0
2
8 10
12 14
Index
Index
(a) 3-bit quantization
(b) 4-bit quantization
Figure 5: Visualization of DT quantization and Linear-2 quantization at b-bit (b = 3, 4) precision.
B
Randomized SVD Method
Given an initial matrix Po E R=x", randomized SVD method computes the eigenvector matrix of matrix A "xx by iterating
= QR
where QR(X) denotes the QR decomposition of matrix X, returning an orthogonal matrix. Since we can initialize Po with the previous result (e.g.. V in Algorithm [ only few iterations are enough obtain accurate estimation in practice. In our experiments, iterate (4) once for Shampoo/CASPR, and iterate (4) twice for K-FAC/AdaBK.
C Quantization Mappings
We present the constructions of different quantization mappings in b-bit quantizers (R in Q). See Figure 5 for the illustration them. Note that T, - (0....... -1)
Dynamic tree (DT) quantization for b-bit quantization maps T, onto (0, U G, where G is set numbers with the following properties: the number G looks like =Qş x 10- where a)
b 2 E + F, where E, F' are integers; b) 9k (px + Pk+1)/2. where k E {O...., 2 1}: For 4-bit quantization, DT quantization maps T4
c) Pj = 0.9j/2f + 0.1, where j E (O...., 2 onto 1-0.8875, -0.6625, -0.4355, -0.2155, -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 T3 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 T4 onto (-1.0000, -0.7511, -0.5378, -0.3600, -0.2178, -0.1111, -0.0400, 0.0000, 0.0044, 0.0400, 0. 0.2178, 0.3600, 0.5378, 0.7511, 1.0000). For 3-bit quantization, Linear-2 quantization maps T3 onto {-1.0000, -0,5102, -0.1837, 0.0000, 0.0204, 0.1837, 0.5102, 1.0000).
D Quantization Error Analyses
We present more quantization error analyses the preconditioners. Recall that we define two kinds of quantization enrors in mapping f of transformation at A E Rmxn (in short errors in f(A) ) of g) in Subsection 3.1 Here we extend them as follows: define the normwise relative emro (NRE) in f of (g1, 02) at A as
IIf(A) [If(A)]le
(A)llr
NRE
and the angle emor (AE) in ff of (91. 92) at as
(f(A), g2 oog [[(A]] Fl92 fogi (A)]l (og](A)]]
AE = arccos
14

---

D.1 Static Analysis
Table 5 an extension Table for Bit=4. Since the diagonal elements of --14 are usually much larger than its non-diagonal elements where A PD matrix, further consider the quantization errors in f(A) A 1/4 Diag( diag(A 1 4-bit precision as shown in Table Table 7 shows
the quantization errors at 8-bit precision.
A large condition number of PD matrix A is indispensable for the superiority of quantizing U over quantizing A, where is the eigenvector matrix of A. We consider contracting the singular value distribution of A = A1 with SVD UDiag(A)U" used in Table 5 by mapping each singular value of A to h(A) = r(> min) + > min- where min the minimum singular value of A and T > the contraction coefficient. Figure shows 4-bit quantization errors in A-1/4 or A 1/4 Diag(diag(A (A44)) of quantizing U A A A= (DD)gg(((())))
Table 5: Quantization errors in f(A) = A-1 of different 4-bit quantization schemes at PD
matrix A. We employ block wise normalization with block size of 64. U is the eigenvector matrix of A and (A)) 1/4 QM quantized matrices and OR orthogonal rectification.
Real-world A A
Synthetic A = A
Mapping R
QM
OR
NRE 4
AE (?)
Mapping R
QM
OR
NRE
AE (
A
X
0.6241
17.319
X
0,4615
17.189
U
X
0.0709
4.0426
U
X
0.1224
7.0144
U
V
0.0455
2.5615
U
V
0.0878
4.9960
DT
B
X
0.0398
2.2802
DT
B
X
0.0853
4.8914
(A, B
X 0.6243
17.364
(A, B
X
0.4649
17.650
U B
X
0.0811
4.6296
U B
X
0.1485
8.5168
(U, B)
0.0604
3,4230
(U B)
0.1224
6,9817
A
X 0.6243
17.293
A
x 0.4465
15,338
U
X
0.0543
3.1066
U
x
0.0942
5.3998
U
J
0.0343
1,9456
U
0.0669
3.8166
Linear-2
B 0.0315
1.8050
Linear-2
0.0661
3.7887
(A, 0.6243
17.301
B
0.4483
15.654
U. 0.0626
3.5833
0.1150
6.5901
U.B) 0.0466
2.6494
0.0941 (U,
5.3716
Table 6: Quantization errors in f(4)= = A 1as/444Daa((iaae (bag(diag(A-1144) of different 4-bit quantization schemes a PD matrix A. We employ block-wise normalization with block size 64. U the eigenvector matrix of A and (g1 (A)) 1/4, QM quantized matrices and OR orthogonal rectification.
Real-world A
Synthetic A A
Mapping R
QM
OR
NRE
AE (*)
Mapping R
QM
OR
NRE
(?)
A
X
0.9549
59.360
A
X
0.6247
25.913
U
X
0.2328
13.287
U
X
0.1994
11.444
U
/
0.1480
8,4365
U
/
0.1427
8.1415
DT
B
X 0.1314
7.5513
DT
B
X
0.1391
7.9813
(A, B)
X 0.9561
59.825
(A, B
X
0.6314
26.948
U B
X 0.2666
15.281
UU B
X
0.2420
13.911
(U. B)
0.1977
11,322
(U, B)
0.1992
11.393
A
X 0.9547
58.336
X 0.6010
20.780
U
X
0.1786
10.213
U
0.1534
8,8027
U
0.1122
6.4096
U
0.1088
6.2176
Linear-2
B
X
0.1041
5.9554
Linear-2
B
X
0.1078
6.1755
(A, B)
0.9548
58.601
A.
B
X 0.6047
21.666
U.
0.2063
11.778
U
B
X 0.1873
10.745
(U. B)
0.1530
8.7337
(U, B)
0.1532
8.7534
15

---

g0.
quanA quan If
150
17.5
qui A
0.8
1.0
quan quen
12.5
9.6
0.3
0.2
7.5
0.4
quan
0.1
-1 Contraction Confiicen
21 -1 Concrccion Coefiident
Contraction Coetrieent T
Contraction Cooffillet
(a) in ((A
1/4
A 1-11-11/1 Errors ((A)
1/4
Diag (diag 1/4))
Figure 6: 4-bit quantization errors in (A) of quantizing U or A at A UDiaa(( We use linear square quantization and orthogonal rectification. The condition number cond(A) maxx AA is around 37235, where A and X are the maximum and minimum singular values of respectively. Contraction coefficients are shown on log, scale.
Table 7: Quantization errors in f(A) of different 8-bit quantization schemes at PD matrix A. where = Ai is derived from the real world as described in Subsection 3.1 employ block-wise normalization with a block size of 256. U eigenvector matrix of A and (A)) -1/4,
QM quantized matrices and OR orthogonal rectification.
((A)
A
S(A) A-1/4 1/4 Ding (ong(dag(44) 11/1)
Mapping R
QM
OR
NRE 4
AE (*)
Mapping R
QM
OR
NRE
AE (*)
A
0.2192
8.3014
X
0.5001
23.644
U
X
0.0060
0.3421
U
X
0.0197
1.1273
U
V
0.0037
0,2140
U
0.0123
0.7022
DT
B
X
0.0029
0.1655
DT
B
X
0.0097
0.5553
[. B
0,2193
8.3051
(A, B
X 0.5003
23.649
U.B)
0.0067
0.3810
U B
X 0.0219
1.2577
U B)
0.0047
0,2712
(U, B)
0.0156
0.8955
A
X
0.2164
7.9751
A
X 0.4875
21.477
U
X
0.0037
0.2121
X
0.0122
0.6994
U
0.0023
0.1312
UUU
0.0076
0.4343
Linear-2
B
X 0.0021
0.1203
Linear-2
B
X
0.0070
0.4035
(A,
B
X 0,2164
7.9755
(A,
B
X
0.4875
21.488
U.B
X
0.0043
0,2439
UUB
X 0.0141
0.8079
U,B)
0.0031
0.1791
(U,
0.0104
0.5935
D.2 Dynamic Analysis
define the normwise relative error (NRE) and angle error (AE) of deviating from A as
NRE B - A
(A, B
AE AE=arccs
|
| B
Consider Shampoo using 4-bit preconditioners for parameter updates, but also recording 32-bit preconditioners at the same time. We extract the left preconditioners L1 and L32 E IR 1200x 1200 of specific model parameter block W 11200x768 every 8000 steps in the Swin-Tiny training on CIFAR- 100 with Adam W+Shampoo. Here L, is a decompressed 4-bit preconditioner, and Lyz is 32-bit preconditioner.
Figure shows the quantization errors during training. For naive 4-bit Shampoo, L -1/4 and L 1/4 computed by Schur-Newton iteration used in Algorithm 4 where 10 -4 For our 4-bit Shampoo, Lg2 is computed by Schur-Newton iteration used Algorithm 4 where e = 10-4 and L 1/4 computed by Algorithm without quantization where = 10-1 1/2 4. We find that 10 for Algorithm used in our main experiments though is effective, yet it can cause a large numerical instability the later stage of training (see Figure
2
16

---

0.6
4-x Shampoo (narr)
1.09
<-bx Shampos (our)
5.55
11 L11 18 8 8
0.6
20.5
bit Stampoo vro
bit Shampos (natre)
ont Sho (nalve) sit Shampoo (our)
4-b1 Shampso (our)
0080
bit Shampoo (our)
0.75
0.2
Step Za (00000)
Step (xabot)
Step (x8000)
Step (40008)
(a) Errors in L deviating from g2
(b) Errors in L 1/4 Diag Diag(diag(Li) L/141)
1/4) deviating
from Lg
1/4 Diag(diag( L522
Figure 7: Quantization errors during Swin-Tiny training on the CIFAR-100 dataset. We use dampening term 10 to compute L4" 1/4 and 11/4
20.9
2166
Stampon Shampoo (oud
Mo.7
Sharrpoo (oug)
6.66
Adit She (naive]
Anr Enre
Belil Shampen (nstue)
8..5
(our)
Shompoo
0.4
203
0.2
6
15
(H80501
(a) Ertors in L 1/4 deviating from Lgz 1/4
(b) Errors in L, 1/4 Diag(diag( L, 1/1) deviating
1/4 Diag(diag( from Lyz
Figure 8:
Quantization errors during Swin-Tiny training the CIFAR-100 dataset. We use dampening term e 10-6 compute - and L L///
E Convergence Analysis
More notations. Given symmetric real matrix A, A means that A is positive semidefinite (PSD), and A > 0 means that A is positive definite (PD). Assume that symmetric matrices A and B and > mean that A
and A - B > respectively. Let A be a PSD matrix and E R. we define -==UUU where UAUT is the Singular Value Decomposition (SVD) A. The Mahalanobis norm of a vector x induced by a PD matrix A || = V T Az. The dual norm of || |IA is denoted by ||- -1 where || = VxTA-1a The spectral norm of matrix A is |||| suP, |||AAx|| |2/||==|| A @ B means the (right) Kronecker product of matrices A and B. vec A) means the vectorization (stacking the rows) of A.
are symmetric, the notations
Algorithm 6 Perturbed Shampoo in the matrix case
Input: Wo E mm Lo 0mxm Rg O Ouxn: Po = 0. po = 0.
1: for /= 1,..,., T do
2: Receive loss function: J. :R:RRmx mm + R
3: Compute gradient: G, (W,)
4: Update preconditioners: J. L1-1 + G.G.;
K. R G.
5: Perturb preconditioners: L, = g(J:); R = g(K))
6: Accumulate errors: Pt = pt-1 + +J: -J--10 Lıllz; 14= 14-1 + K; Rall
7: Update parameters: W++1= Wt W,-n((s P.) L.) 101 H.) + -R)
We consider quantization as perturbation and present the perturbed Shampoo in Algorithm 6 for convergence analysis. The regret bound of the perturbed Shampoo can be found in Theorem Complete proofs can be found in Appendix F We first introduce some basic technical tools, and the details of them are in [8 2T).
Lemma 3. Let A, A', B, B' be matrices of appropriate dimensions, and u, be two column vectors. The following properties hold:
(1) (A & (A' & B') (AA') ② (BB);
17

---

(2) (A & = (A & ):
(3) If B = 0 and R, then (A B) (A' @ B");
(4) If X A' and B = B', then A % B A' Ø B';
(5) tr(AB) tr tr()))rrBB:
(6) vec(uv = u X
Lemma 4, Let G E mRmxx, E Rmxm R E Raxn, then holds that
(L R ved vec vec (LGR).
Lemma 5. Assume that 0 X, x Y; for =....... l..... Assume further that all X; commute with each other and all Y, commute with each other. Let 01 ... an 2 such that Σi=1 0 1, then
X ... X 3 Y Y....
Lemma 6. Let S a S 1 and 0 3 X Y, then X 1033 Ya Lemma's 7. Let A 0 and B > O, then it holds that A = B if and only if B B A-1
Lemma 8 (von Neumann). Let A, B Rmxx and q = min (m, n). Let σi (A) 2 .... og (A) and T1 (B) 2 ... (B) denote the non-increasingly ordered singular values of A and B, respectively. Then
(A.B) S E (B)
Lemma Assume that function ft is continuously differentiable and convex on Rdd, and matrix R RR,1 > 0, define W wt nH, gt. where
H. > 0 for 1,...,T, Given
(w,... Then for any w0 E Rd, we have
T
- f((w)) 5
2n E0 (1w w" P.
1|w-11 w0 11 +
> (1leelli,))
2 t=1
Lemma 10. Let Letgi..... g1...., 9T be a sequence of vectors. For > O, define H, = (pl E.= 1 9 Then we have
T
Σ
H. (lgel )2 2tr(Hr))
Lemma 11. Assume that GI...., Gr E T mxx are matrices of rank at most r. Let s fort 1.... T. Then for any 2 0,
T
Imn
Σ gigt (cIm + G.GG & (cI, + GG.
t=1
The key to the convergence proof of Algorithm
6 forming a PD matrix sequence ( (H, [H;)T=, which satisfies 0 < H, I Hr. To achieve we gives the following lemma extended from Lemma in the Appendix of 40
Lemma 12. Let {X,)=I be sequence of symmetric matrices, and Ar Σt X. where
Mt=T
t= I. I....,T. Suppose we have two sequences of symmetric matrices (Y) |=T (Ze) sequence real numbers t=) satisfying
t=t and
Y. Z,-1 + Xt, P: == - I|Y, Z,12, Zo = 0, Pu = 0.
Define B. I + Z₁ where I denotes the identity matrix. Then for 1,..., T, we have
B. Bi-1 +X, A. - B, - 2p. A..
Theorem 1. Assume that the gradients GI,... GT E mm are matrices of rank at most Then for any W E Rmmxn and e > 0, if = D/V2r, the regret of Algorithum 6 is bounded as follows,
T
T
Σ f(())
1/4 + tr 1222/ + tr(Rtt/ Σ f((W) < 2 mpt
E
where maxtelr
W - F, L. cl + Σr_ G.G., and RG cl, + Σt G.G.,
Though we get a convergence guarantee of Algorithm 6 the upper bound given by Theorem I very slack, since 2 1/4 is about the same tr(2)/" for 4-bit quantization schemes in practice.
18

---

F Proofs
Lemma 1. Let A be PD matrix whose SVD UAU,, where U = [u.]] an orthogonal matrix
and A =diag((A,))) is diagonal matrix. Given a perturbation AU = [Au,] and IR, we define
B: B:=(UAU1) and AB ((U=
" BB
(1) +UU+AU is orthogonal and there exists a R such that ||Au; l2 S a, then
|
[B S 2a.
(2) UU++AU is orthogonal and there exists such that (u,, u, + Aui) 11-3 2 then
(B, B B
B AB 2 (1-3)².
B F
Proof. (1) Since U and U AU are orthogonal, we have
B = UUUAU B+ (U AU)AA(U AU)
by definition. This leads to
AB U + AUA + +U)).
The Frobenius norm satisfies the triangle inequality and is orthogonality invariant. Hence,
AB F U + AU + AU)
AU F + |AU (U + AU) ||F
IA r + A F = AU ||
2 Σ |A = 2v Σ X22 Au.
< 2
Σ, Ajtta 2ay Σ 2a A
20 BlI
(2) Similar to (1), we have
AB AA AUAU + AU
From (u,, u. + A 2 3 O, we get 0 2 (u,, Au,) 2 -B 1 because
u. ll; + Au:lla 2 (ui, u; Au.) = 1+ (ui, Au.) 2 B1-320,
holds due to the orthogonality of U and U AU. Hence,
(B, AB) tr(2U A 2 T +UA +UAUUAUAAA AUA-AUT)
Σ 2Af + tr (E,A,aa)
( A;Au,
(A2 Au.)) + AAsg(ussAu))
2 24, Aui)) + [ (u,, Aui)
Σ AP + (u,, Au))? -1]
E, XF*((1 E,7a((1 - 3)2 1] [(1 [3)= 1| A 11
[[1 8)3 B|]
(11 (1- (B, B).
Therefore, we have
B, B + AB
BB B - AB)
= 1 BB AB)
2 (1 -322.
B B + AB
(B.B)
B, B)
The proof is completed.
19

---

Lemma 2. Let A be PD matrix of order m++n whose SVD is UAUT, where m. n. E N+. n = Im,
= [[ an orthogonal matrix and A diag([A.)") is a diagonal matrix. Assume that miag((c ). e Z L, and À > 0. Given a perturbation AA diag((0 A )
and s R. we define B := y and = (U++UUUA+AAJT))
(1) If AAxi = (k - 1)A1ax1 where > 0, then
|AB|| B
VIk h1 (s.1). 2 +1
Moreover, h1 (s,1)) decreases monotonically with over (-20, 0) and increases monotonically with l over (0, +20). (2) If AAxx = (tc 1)x1nx1 where t >> O, then
(B, B
11 +
= h2(1)-
B B B AB
V(L+ + c²a)
Moreover, (1) decreases monotonically with over (0, (c/t)*] and increases monotonically with l over ((c/t)" +00).
(3) If AAnxi = (tc 1)Alnxl where to > 0 and I = (c/t)², then
| B
1
(B. B AB)
+K+
B B + AB |F
V ks 1/k²
Proof. (1) Since U is orthogonal, we have
AB|E
[[ + AA)* A V 1|12, |B] |A Vmc2 nd.
Hence,
|AB|I Vnlk* - B
Mik 1|
(s, 1) 2 0.
Vinc2 + Vo25 +1
easy to check that h increases monotonically with over (0, +oo). To prove h1 decreases monotonically with over (-o0, 0), define
(s) 1 (hn( ②² +
(k* (k--13 1)2
Consider the derivative of 91
01(s) (e2x + 1)2(k= - 1)k" In (k* 11)22222 Ine
(c20 + 1)2
2(k 1) (125 + l)k In k -1+ 1) Inc)
(c2s + 1)2
If < 0 and 1, then 1s < 0, k- > leading to o (s) < since e 2 0; Similarly, if s <
and 0 < & < 1, then kt 1 2 0, At Ink < 0 leading to g (3) < 0. Thus 91 (s) is monotonically decreasing function for s < 0, which implies that h2 decreases monotonically with over (-o0, 0).
(2) Similar to (1), we have
| = V 2 + nA
IB+ABI = Vntza + >.
Besides,
BB AB) (UAA AAYUT)= tr AA)*) (mc2² + nnc
Hence, we get
B, B+AB
nt + mc
11a c
V(1+ (1++)2)((1 h2(1) 2 0.
B F
B+
AB
F V(m + nn22)(n + mc20)
20

---

To prove h2 decreases monotonically with 1 over (0, (c/t)* and increases monotonically with over ((c/t)", +00), we define
(It" + c2)2
(2) = (h2(1))² (1 + (12)((1 c20)²
whose monotonicity is equivalent that of h2 for 1 > 0. Consider the derivative of 92
2²02 + 20 2² ) (t' tf- (c" 22c2
(1((1) = 120 +1+ 22
(t=012 +1 f20 c22)
If s or te then (2( = 1. If + and to # 1, then (ta - f2 cO)² 0, (ce - t2c200 >
In this case, let (2) = 0, we get
(1 111-te = (1- -
which implies that / = (c/t)². It is easy to see that decreases monotonically with over (0, (c/L)
and increases monotonically with I over ((e/t)*, +00).
(3) According to (1)(2,, we can easily get
|B
F
2 | AB V2+ + 1/k2
B
B
v
+I
The proof is completed.
Proposition 1. Let A be PD matrix of orler m-+n whose SVD is UAUT, where m, n E N
Im, U [uu]] an orthogonal matrix, A diag([cA1] mx 1m x1 xx1 [T). c 1000, and
Given AU = [Au,], AA diag [O] mx A [T), and 8 $ -0.25, we define
(UAU
B := (UU + AU + )s, and BN := (U + AA U + AU is orthogonal,
|Au, S 0.1, (u. Au;) 0.005, AAax (0.02c 1)A1nxI, and 1= (c/0.02)*, then
2 B1 Bll
B2 B
6 6(1
(1)
B1)
50.0 (1
( B₂)
50.45
B
| BI
B F
B2
Proof. According to Lemma we have
B
B
B, B
B S 0.2,
2 (1 0.005) > 0.99.
B B
On the other hand, from Lemma
we get
B2
BlI
B2)
2
fi(z),
v2+x+1/5 f2(x),
B
F
Vr
B B2llF
where = (0.02c)* E (0,20 (0,20-1/4). It is easy to verify that fi decreases monotonically and fz increases monotonically for 0 x < Hence
f(x) (20 0.4, fa(x) S f2(20 < 0.94.
The proof is completed.
Lemma 12. Let {XiM=TT be sequence of symmetric matrices. and A, E X, where
== 1, I.....T. Suppose we have two sequences of symmetric matrices {Y;)=T [Z sequence real numbers {=e)== t-= satisfying
T- and
Y. = Zi-1 + X,, = P + I Z1122,
Zo = 0, Pu = 0.
Define B, = l.l + Z. where I denotes the identity matrix. Then for = I,..., T, we have
B, B X1, A B, = 2p, I + A..
21

---

Proof.
Note that for any symmetric matrix S, it holds that || ł S. Then we have
(pt - - + Ze = 1 Zr|12I Z = Y.
Adding Pr II on both sides, we get
B, - p + Z = P1 Y, = P1- Z + X, = B + X
Hence
B. Σ (B, BB-1) Σ X, A.,
5=1
On the other hand, we have
Z: 3 Zz Yill21 + Y, = (pe Pt-1) I Yt.
Adding l on both sides, we get
B, Pc + Ze 3 (2pt - Pt-1) + Yg
2(p: Pt-1) + Pt | + Z,-1 + Xt
B + +(: Pt- + X..
Hence
B D(B. BB-1) 3 E2(p. Po-1) + Σx, 2p,I + A..
The proof is completed.
Theorem 1. Assume that the gradients Gt...., GT E gm are matrices of rank at most r. Then for any W E Rmxx and > 0, if = D/v2r, the regret of Algorithm is bounded as follows,
T Σ I(W.)
T
Σ ((W') S D2r 14 mpT tr (1 )12 mpp tr (r( ((
1/4
where D [ W. W' L, eIm + Ets G.G. and R. cI, + Et G.
Proof. Define L, = + Pt Im + L,, R = (e + μ4 In + R. According to Lemma 12 L, and R. positive definite. Recall the update performed in Algorithm
W = W. - "n_ 1/4 G
For 0, let H. //44 @ R/ g: vec(G.) and wi = ved Wi). Due to Lemma 3 and Lemma we have
u++1 = Wt H,, 9t.
Lemma 12 implies 0 x L, 1 1 Lr, 0- x RS 1 ... - Rr. Thus, according to Lemma and Lemma
4)
we get
0 Hi
H
Let Hg 0. By invoking Lemma 9 and Lemma we obtain the regret bound
(W,) (( S
wi w (HH H. H.-1)(w: w²)
" (1g (r.) )
27
D
T
In H. - H.-1) + " T
<
> (1ggeli,.)) )
21
2
D0
tr
H +
Σ
(llgrl) (r,, )a.
2m
2 t=l
22

---

where D = | | ---w w2112 [7) [ W P and w" vec (W.).
Define H, (reI+ + E. 0.9.)!/ Lemma 11 and Lemma 12 imply that
H. 3 VrL. & R I VrH,
Using Lemma 7 and Lemma 10 along with the above equation, we obtain
(llg: li. )² S (1lge i
>= $ 2 s 2rtr(HT).
1=1
Consequently, using Lemma [5] and Lemma
12
we get the desired regret bound
D2 Σ f(W.) Σ ((W)) v
tr( HT) V2r DY )tr( (12
+ m
2n
s / 1/4 tr tr(2) + tr (R(
by choosing 17 = D/ The proof is completed.
G Experimental Details
We use one RTX3060TI GPU under the PyTorch 2.0.IL framework for DNN training on the CIFAR-100 and Tiny-ImageNer datasets, use one A800 GPU under the PyTorch 2.0. 10..+CUDA11.7 framework for DNN training on the ImageNet-1k Ik and C4 datasets, and use two NVIDIA L40S GPUs under the PyTorch 2.0.1 A.UDIIS framework for DNN training the OWT dataset. To obtain the total peak memory consumption per GPU, we call "torch.cuda.max_eent wtw.daaaaaammararraa.aa www.ommmllloatttededdeeo
We set ViT-Base/32 on the ImageNet- Ik 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/ 22/4 Consequently, the memory savings are approximately 7 times, calculated as 32 (4++0.5) In the future, we may adopt double quantization [ to further reduce memory consumption.
wwr.eknnss.cudn...o "enndnnbbnccmarr""" to "False" for all the experiments, except when training
For SGDM, Adagrad or A 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,
set the momentum to 0.9 and use initial learning rate of 0.1. For Adagrad, we set and use an initial learning rate of 0.01. For AdamW, we set 31 = 0.9, 32 = 0.999, and 10 and use initial learning rate of 0.001. For quantization settings, we employ block-wist 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, f 1 = by default. Shampoo and CASPR precondition blocks from large matrices and the maximum order of 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 the CIFAR- 100 and Tiny-ImageNet datasets.
e = 10 10-10 10
10-6 B = 0.95 and
Settings on training CNNs on CIFAR- or decay is 0.0005. Data augmentation includes random crop and horizontal flip. For Shampoo, we set Ti = 100 and T3 = 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).
Minibatch size is set to 128. Weight
Settings on training transformers on CIFAR- 100 or Tiny- ImageNmt. We set patch size of for ViT-small on the CIFAR-100 dataset, and patch size of 8 for ViT-small on the ImageNt dataset. For training Swin-Tiny on the CIFAR-100 dataset, we use a patch size of and window size of 4. For training Swin-Tiny the Tiny-ImageNet dataset, we use patch size of 4 and window
23

---

size of 7. Minibatch size is set to 128. We run Adagrad/Adam Adagrad/ Alagrad/Adam/ii AdamAddaaWaa.a Shhampo for 100 epochs. Weight decay is 0.0005 for Adagrad, and is 0.05 for Adam ClamW/NadammW We use the cosine learning rate schedule. Data augmentation follows the source code 25 For Shampoo, we set 71 = 100 and 72 500. With the exception of certain optimizer settings, the configurations used for ablation studies are identical to those outlined above.
for 150 epochs and
Settings on training ResNet50 on ImageNet-Ik. We run SGDM for 120 epochs and SGDM+Shampoo for 100 epochs. Minibatch size is set to 256. Weight decay is 0.0001. We adopt the multi-step learning rate schedule (the learning rate multiplied by 0.1 for every 30% epochs with a linear warmup at the first 5 epochs). Data augmentation includes random resized crop, horizontal flip, and color jitter. For Shampoo, we set Ti 200 and T2- 1000.
Settings on training ViT-Base/32 on ImageNet-Ik, We run A for 150 epochs and AdamW+Shampn for 120 epochs. Minibatch size is set 512. Weight decay is 0.05. We the cosine learning rate schedule. Data augmentation follows the configuration for training ViT-Base/16 in [44), excluding repeated augmentation. For Shampoo, we set 200 and 72 1000.
Settings on training GPT-2 on OWT. We run AdamW with 10% warmup steps. Total batch size is set to 480. Batch size is set to 24 for training 124M GPT-2. Dtype is bfloat16. Weight decay is 0.1. For Shampoo, we set Tj = 200 and T2 = 200. For our 4-bit Shampoo, we use Schur-Newton iteration used in Algorithm 4 to compute the inverse root of preconditioner for training stability.
Settings on training LLAA MA-2 on C4. We run Adam with 10% warmup steps. Total batch size set to 512. Batch size is set to 256 for training 130M LLAMA-2 and is set to 128 for training 350M LLAMA-2. Dtype is 66. Weight decay is O. For Shampoo, we set
200 and T2 200.
Settings on K-FAC and AdaBK. K-FAC/AdaBK preconditions layers without limiting the size of a preconditioner. We set 0.9, Ti 200, and T2 2000. We use e 0.1 for K-FAC and e 0.001 for AdaBK. For 4-bit K-FAC/AdaBK, we set t = 0 and l2 = 0 (i.e., no orthogonal rectification).
Settings on schedule free optimization. We use the code from 6 to train ResNet34 with SGDScheduleFree and Swin-Tiny with AdamWScheduleFree, For SGDScheduleFree, set 1r=1.0,
and warmup_stern steeuppsteps=2000. For AdamWSSchedullFFre, we set 1r=0.0025, weight_decay=O.nl and warmup_ harmup_steps=10000
Settings on M-FAC. We use the code from [15] and set ngrads=32. damp=0.1, The other hyperpa- rameter settings of M-FAC is the same as that of SGDM used for ResNet34 training.
H
Additional Results
H.1 Image Classification
More learning rate schedulers. Table Table&&shore shows the performance and wall-clock time of training ResNet34 on CIFAR- 100 with cosine learning rate decay. By comparison, SGDM+Shampoo still converges faster than SGDM, and have slightly better test performance.
Table 8: Performance and wall-clock time of training ResNet34 the CIFAR-100 dataset with cosine learning rate decay. TA = test accuracy, and WCT = wall-clock time.
Epochs
Optimizer
TA (%)
WCT (min)
200
SGDM SGDM
79.67
116.0
300
79.83
172.7
200
SGDM + 3it Shampoo
80.39
152.7
200
SGDM + 4-bit Shampoo (our)
80.22
161.7
We also provide the results of training ResNet34 and Swin-Tiny on CIFAR-100 with schedule- free approach [6] in T From it one can see that Adam WScheduleFree achieves comparable performance to Adam W with cosine decay, while SGDScheduleFree underperforms compared to SGDM. We observe that this schedule-fres algorithm shows rapid improvements in training and test accuracy during the early training stages, but may fail to achieve higher test accuracy ultimately (see Figure 9. Anyway way, these methods are still worse than our AdamW+44 Shampoo.
24

---

Table 9: Performance and wall-clock time of training on the CIFAR-100 dataset with cosine learning rate decay and schedule-free approach, ResNet34 trained for 300 epochs and Swin-Tiny trained for 150 epochs. TA = test accuracy, and WCT = wall-clock time.
Model
Optimizer
TA (%)
WCT (min)
SGDM
79.83
172.7
ResNet34
SGDScheduleFree
75.63
169.6
AdamW AdenWScche
76.69
318.6
Swin-Tim Tiny
76.58
321.9
ResNet34 CIFAR-100
Swin-Tiny CIFAR-100
75
Test Accuracy (%)
70:
Accuracy
70-
65
65
60-
SGDM
55
Adam
SGDScheduleFree
AdemWScheduleFree
50
100
150
200
250
300
20
40
60
80
100
120
140
Epoch
Epoch
Figure 9: Visualization of test accuracies on the CIFAR-100 dataset with cosine learning rate decay and schedule-free approach.
More optimizers. Table TO shows results of training Swin-Tiny on CIFAR-100 with NadamW. Adagrad and Adagrad+Shampoo. One that 4d Shampoo converges faster than Adagrad with ignorable extra memory overhead, and also has higher test accuracy. Besides, though WadamW I is slightly better than A, it still worse than our Adam Wit-bit Shampoo.
Table 10: Performance, wall-clock time, and memory cost of training Swin-Tiny on the CIFAR-100 dataset. TA = test accuracy, WCT wall-clock time, and TMC = total GPU memory cost.
Optimizer
TA (%)
WCT (min)
TMC (MB)
NadamW
77.11
342.4
1465.8
Adam W 3it-bit Shampoo
79.34
260.8
2036.0
AdamW + 4-bit Shampoo (our)
78.63
273.3
1543.9
Adagrad
66.56
294.6
1354.9
Adagrad 32-bit Shampoo
73.55
245.3
1930.4
Adagrad 4-bit Shampoo (our)
72.66
259.6
1433.0
M-FAC [15] is matrix-free method computing inverse-Hessian vector products with many gradient copies. is not memory-efficient for M-FAC maintain m dense gradient copies (m 1024 in its official code). Table 11 shows that both SGDM+32-5P 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 rank-one recond-orde optimizer and is memory 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 422 The test accuracies of SGDM+EVA and SGDM+Shampoo are 73% and 74.5%, respectively.
25

---

Table 11:
Performance and memory cost of training ResNet34 on the CIFAR-100 dataset with cosine learning rate decay. All the optimizers are run for 200 epochs. TA test accuracy, and TMC = total GPU memory cost.
Optimizer
SGDM M-FAC (m==32 SGDM + 32-bit Shampoo SGDM 4-bit Shampoo (our)
TA (%) 79.67
78.56
80.39
80.22
TMC (MB)
822.03
3424.8
1441.8
908.4
Table 12:
Performance, wall-clock time, and memory usage per GPU on natural language modeling
tasks. VI. validation loss, WCT wall-clock time, and TMC total GPU memory cost.
Dataset
Model
Optimizer
VI.
WCT (min)
TMC (MB)
AdamW
3.214
346.9
47026
AdamW 32-bit Shampoo
3.184
353.7
48813
LLAMA-130M
AdamW 4-bit Shampoo (naive)
3.200
353.5
47316
AdamW 4-hit Shampoo (our)
3.194
353.1
47318
C4
AdamW
2.939
2687
54184
AdamW + 32-biit Shampoo
2.908
59149
LLAMA-350M
AdamW + 4-bit Shampoo (naive)
2,930
2753
54894
AdamW 4-bit Shampoo (our)
2.924
2795
54894
AdamW AdamW + 3--2-bit Shampoo
2.954
2310
27010
2.936
2330
28490
OWT
GPT2-124M
Adam W + 4-bit Shampoo (naive)
2.953
2359
27209
Adam 4-bit Shampoo (our)
2.944
2311
27209
3.24
LLAMA-130M on C4
2.90
GPT2-124M on OWT
3.23
AdamW+32-b6 -3-------ll bit Shampoo
Shampoo (naiive)
2.974
Adami++4 Shampoo (naive)
AdamiW 32-bit Shampoo
3.22
AdamW+4-bil Shampoo (our)
AdamW +4 b-tb Shampoo (our)
2.96
Valldation
3.21
Validation
2.95
3.20
3.19
2.94
3.18
260
360
2100
2200
Wall-clock [min)
Wall-clack Time (min)
Figure 10:
Visualization of validation loss on the C4 and OWT datasets.
H.2
Natural Language Modeling
Models, datasets, and hyperparameters. We train 124M GPT-2 32 for 60k steps 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 the C4 dataset (33 with one A800 GPU. See Appendix G for experimental details.
following 43
Main results. We show the performance, wall-clock time, and memory cost in Table 12 and the validation loss curves in Figure [0] As with the vision tasks, our AdamW++4 Shampoo consistently outperformed AdamW and naive AdamW+4-bi Shampoo in terms of performance, and AdamW+32-br Shampoo in terms of memory usage.
Memory efficiency. We further check the memory usage by increasing token batch size for language model, which calculated as the batch size multiplied by the context length (see (43). To train LLAMA2-7B on the C4 dataset using single A800 GPU (with maximum memory of 81,920
26

---

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 preconditioner for training LLA.MA27B 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 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.
Table 13: Memory cost of training LLAMA2-7B on the C4 dataset with different optimizers. One A800 GPU with maximum memory of 81,920 MB is enabled. TMC total GPU memory cost, and OOM out of memory.
Optimizer
Batch Size
TMC (MB)
8-bit Adam W
64
60135
8-bit Adam W
128
68689
8-bit Adam
256
OOM
8-bit Adam W + 32-bit Shampoo
2
OOM
8-bit AdamW 4-bit Shampoo (our)
64
74561
8-bit AdamW 4-bit Shampoo (our)
128
OOM
27

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope?
Answer: [Yes]
Justification: We propose the first second-order optimizers with 4-bit states by taking Shampoo as an example, while preserving the performance achieved with 32-bit optimizer states.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contribu- tions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine include aspirational goals as motivation as long as it is clear that these goals not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We discuss the limitations of the work at the end of the paper.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to vi- olations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by review- ers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't It acknowledged the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: We provide all the proofs in the Appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross-referencee
All assumptions should be clearly stated or referenced in the statement of any theorems.
28

---

The proof' 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 supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experi- mental 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: [Ys]]
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, 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 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 novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility results, access to a hosted model (c.g., in the case a large language model), releasing of a model checkpoint, other means that are appropriate the research performed.
also be provided via detailed instructions for how to replicate the
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. For example
(a) If the contribution primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution a new model (c.g., a large language model), then there should either be a way to access this model for reproducing the results way reproduce the model (e.g., 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, may be that access 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.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, 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 (nttps://nips. guides Sttlee//ooodduuubiissssulolooooo 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 (c.g., for new open-source benchmark).
29

---

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https:// for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide seripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
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.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: Error bars Guidelines:
reported because it would be too computationally expensive.
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars. confidence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1 1-sigma error bars, but one should state it. The authors should preferably report 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed reproduce the experi- ments?
Answer: [Yes]
Justification: We present the implementation details of all the experiments in the main paper and Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
30

---

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experi- mental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make it into the paper).
9. Code Of Ethics Question:
Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification:
We review the NeurlPS Code of Ethics and paper conforms it.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No. they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make
preserve anonymity (e.g., if there is special consideration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive socictal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: We discuss both potential positive societal impacts and negative societal impacts 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., deploy- ment 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 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 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 unintentional) misuse of the technology.
If there 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 system learns from feedback over time, improving the efficiency and accessibility of ML..
11. Safeguards
Question: Does the paper describe safeguards that have been put place for responsible release of data models that have high risk for misuse (c.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification: We present 4-bit Shampoo for memory efficient training of deep models. poses no such risks.
Guidelines:
The answer NA means that the paper poses such risks.
31

---

Released models that have 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 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 the package should be provided. For popular datasets, paperswithcode com/ 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.
13. 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 datasct/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.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsoureing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: Our paper does not involve crowdsourcing nor research with human subjects. Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code Ethics, workers involved in data collection, curation, or other labor should be paid least the minimum wage in the country of the data collector.
32

---

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Sub- jects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects. Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may required for any human subjects research. you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS 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.
33


Reviews:
[{"review_id": "cDSr2QpZlf", "reviewer": "Reviewer_b2wM", "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.", "questions": "Please refer to the weakness.", "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.", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- This study provides both empirical and theoretical justification for the quantization approach, demonstrating that quantizing the eigenvector matrix of a preconditioner significantly reduces quantization error.\n- The experiments span various datasets and models, consistently achieving performance that matches that of the full-precision Shampoo.\n- Both the speed and memory costs are evaluated in the experiments.", "weaknesses": "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the insightful and positive comments. In the following, we provide our point-by-point response and hope our response helps address your concerns. We also look forward to the subsequent discussion which further helps to solve the current issues.\n\n**1)  Breakdown of memory usage.**\nThanks. Per your suggestion, in the table below, we have reported the additional memory cost incurred by injecting the vanilla 32-bit Shampoo and our proposed 4-bit Shampoo into the vanilla optimizer. The extra memory cost primarily comes from the preconditioning matrices L and R, and their inverse roots. See Eqn. (1). \n\nTo compute this extra memory cost, we use the data from Table 2 of manuscript, and subtract the total memory cost (TMC) of the vanilla optimizer \"A\" from the TMC of \"A + 32/4-bit Shampoo\", where A can be SGDM and AdamW. Comparatively, the extra memory cost of the 32-bit Shampoo is seven times greater than that of our 4-bit Shampoo, demonstrating the superior efficiency of our 4-bit Shampoo. This high efficiency is also analyzed and emphasized in Appendix F.\n\n\n|   Dataset   |    Model    |          Optimizer          | Memory Cost |\n| :---------: | :---------: | :-------------------------: | :---------: |\n|  CIFAR-100  |  ResNet34   |    SGDM + 32-bit Shampoo    |  619.8 MB   |\n|  CIFAR-100  |  ResNet34   | SGDM + 4-bit Shampoo (our)  |  86.37 MB   |\n|  CIFAR-100  |  ViT-Small  |   AdamW + 32-bit Shampoo    |  532.0 MB   |\n|  CIFAR-100  |  ViT-Small  | AdamW + 4-bit Shampoo (our) |  71.70 MB   |\n| ImageNet-1k |  ResNet50   |    SGDM + 32-bit Shampoo    |  630.2 MB   |\n| ImageNet-1k |  ResNet50   | SGDM + 4-bit Shampoo (our)  |  89.03 MB   |\n| ImageNet-1k | ViT-Base/32 |   AdamW + 32-bit Shampoo    |   1534 MB   |\n| ImageNet-1k | ViT-Base/32 | AdamW + 4-bit Shampoo (our) |  204.1 MB   |", "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": "RwrVTlKOc4", "reviewer": "Reviewer_e8nk", "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.", "questions": "Please refer to the section of weaknesses.", "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.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. This paper introduces an innovative memory reduction technique using 4-bit quantization for second-order optimizers.\n2. This paper provides a comprehensive evaluation across various neural network architectures and datasets.\n3. This paper includes practical implementation details with a planned release of source code for accessibility.", "weaknesses": "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.\n2. **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.\n3. **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.\n4. **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.\n5. **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. \n6. **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.\n\n[1] Duchi J, Hazan E, Singer Y. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 2011, 12(7).\n\n[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.\n\n[3] Lin Zhang, Shaohuai Shi, Bo Li, Eva: Practical Second-order Optimization with Kronecker-vectorized Approximation, ICML 2023", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the insightful and valuable comments. \n\n**1) Evaluation scope.**\nThanks. We trained medium-sized language models, including 124M GPT-2 and LLAMA-2 130M/350M. The results are reported in Table 2 and Figure 1 of the rebuttal PDF, which are consistent with those of vision tasks. Due to the complexity and limited rebuttal time, we were unable to conduct reinforcement learning experiments. Since our algorithm design is not task-specific, we believe its improvements should be consistently achievable on other tasks as well.\n\n**2) Time cost introduced by quantization and dequantization.**\nThanks. Block-wise quantization and dequantization are indeed very fast, since they can use the cuda-kernel level parallelism. As a result, they are widely adopted in activation compressed training like [1] and memory efficient optimizers like [2].  As shown in Figure 1(a), the quantization and dequantization in AdamW+4-bit Shampoo (naive) only bring around 3% extra overhead, while greatly improving memory efficiency by 24%. \n\n**3) Quantization error analysis over training.**\nThanks. Quantization-based approaches always try to reduce the quantization errors compared to their full-precision counterparts, so as to expect similar performance. Quantization errors will accumulate along with the training, but they are controllable. As evident from the training curves in Figures 1 and 4, no jump occurs. We experimentally analyzed the impact of quantization errors on the training process (see Appendix D2). The quantization errors of our 4-bit Shampoo are lower than those of naive Shampoo in most of cases and these errors are all bounded. \n\n**4) Time cost introduced by orthogonality rectification.**\nThanks. We follow the vanilla Shampoo method and perform orthogonality rectification (OR) every $T_1$ or $T_2$ iterations ($T_1$ = 100/200, $T_2$ = 500/1000). See line 5 and line 9 in Algorithm 3. This lazy update strategy greatly reduces the computational cost. We trained ResNet34 on CIFAR-100 with 4-bit Shampoo, both with and without OR, resulting in training times of 161.7 and 161.0 minutes, respectively, with OR introducing less than 1% extra overhead. \n\n**5) Comparison with Adagrad, M-FAC, and EVA.**\nThanks.  For Adagrad, it is not as widely used as AdamW due to its inferior performance (e.g. see Table 7 in [2]). We integrate Shampoo with Adagrad, and run Adagrad for 150 epochs and Adagrad+Shampoo for 100 epochs. The table below shows that when training Swin-Tiny on CIFAR-100,  Adagrad + 4 bit Shampoo converges faster than Adagrad with ignorable memory overhead, and also has higher test accuracy.\n\n|           Optimizer           | TA (%) | WCT (min) | TMC (MB) |\n| :---------------------------: | :----: | :-------: | :------: |\n|            Adagrad            | 66.56  |   294.6   |  1354.9  |\n|   Adagrad + 32 bit Shampoo    | 73.55  |   245.3   |  1930.4  |\n| Adagrad + 4 bit Shampoo (our) | 72.66  |   259.6   |  1433.0  |\n\nFor M-FAC, the following table shows that our SGDM+4bit Shampoo enjoys much higher efficiency than M-FAC (m=32) on ResNet34, and enjoys higher test accuracy. This is because M-FAC needs to maintain m dense gradient copies (m=1024 in its official code), and is not memory-efficient.  Here we run all the optimizers for 200 epochs. One can observe that even M-FAC using m=32 already requires much more GPU memory, let alone M-FAC using m=1024, and also suffers from worse performance than SGDM and SGDM+32/4 bit Shampoo. \n\n| Optimizer |   SGDM    |   M-FAC   | SGDM +32 bit Shampoo | SGDM +4bit Shampoo (our) |\n| :-------: | :-------: | :-------: | :------------------: | :----------------------: |\n|  Memory   | 822.03 MB | 3424.8 MB |      1441.8 MB       |        908.40 MB         |\n| Accuracy  |  79.67%   |  78.56%   |        80.39%        |          80.22%          |\n\nRegarding EVA, it is a rank-one second-order optimizer and is memory-efficient. We trained ResNet34 on CIFAR-100 with SGDM+EVA, but despite extensive hyper-parameter tuning, we failed to achieve acceleration over SGDM. Instead, we cited EVA's result of training VGG-19 on CIFAR-100 for 200 epochs (see Table 2 in [5]). The test accuracies of SGDM+EVA and SGDM+Shampoo are 73% and 74.5%, respectively. \n \n**6) Theoretical analysis, e.g., convergence guarantees \\& mathematical justification for quantization.**\n\nRegarding convergence, we do not provide an explicit proof in this work, as our primary focus is on the practical application of optimizers to reduce GPU memory usage. This approach aligns with many other works on optimizers that prioritize practical solutions over theoretical analysis, such as [2][4]. \n\nHowever, we find that the proof techniques of Theorem 1 in [3] can be adapted to demonstrate the convergence of our optimizer. Specifically, we can generalize Lemma 2 used in proving Theorem 1 in [3] and then use this extended lemma to prove the convergence of quantized Shampoo. For further details, please refer to  Lemma 1 in the \"global” response. Note that EVA does not prove convergence, but rather demonstrates that its update step is more aggressive than that of K-FAC from the trust region optimization perspective.\n\n \n\nIn terms of quantization, we do provide a mathematical justification in Section 4 of our manuscript. We observe that the eigenvalues of Shampoo's preconditioner decrease very rapidly. So we assume that the eigenvalues follow a certain distribution, and theoretically analyze why quantizing the eigenvector matrix of the preconditioner outperforms direct preconditioner quantization. \n\n[1] Division: memory efficient training via dual activation precision. ICML, 2023.\n\n[2] 8-bit optimizers via block-wise quantization. ICLR, 2022.\n\n[3] Block low-rank preconditioner with shared basis for stochastic optimization. NeurIPS, 2023.\n\n[4] Adafactor: Adaptive learning rates with sublinear memory cost. ICML, 2018.\n\n[5] Eva: Practical second-order optimization with Kronecker-vectorized approximation. ICLR, 2023.", "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": "ViWqjiyII9", "reviewer": "Reviewer_LFix", "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.\n\nThe 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.", "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.\n2. Please add a short paragraph describing DT quantization, to make the presentation self-contained.\n3. A graph from x \\in [-1,1] to  Q(x) would be nice for each of the quantization schemes.\n4. 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.\n5. \"Singular values are scaled by log 10.\" --> \"Singular values are shown on a log_{10} scale.\" would be a lot less confusing.", "limitations": "N/A", "rating": 7, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "The paper focuses on one aspect of second order optimization, and analyzes it thoroughly. The theoretical justification of the method that the authors provide is intended to provide evidence that this method is better than the naive method. They provide a reasonable set of experiments justifying their work. In addition, the method has the potential of being immediately used in practice.", "weaknesses": "The method should help in optimization of large networks, but no results are presented on large networks. \nThe presentation has some minor issues, noted below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the insightful and positive comments. In the following, we provide our point-by-point response and hope our response helps address your concerns. We also look forward to the subsequent discussion which further helps to solve the current issues.\n\n**1) Training of large-scale models.**\nDue to the limited computing resources and rebuttal period, we cannot afford sufficient pre-training of very large models. Instead, we train medium-sized language models, including  124M GPT-2 for 20k steps on the OpenWebText (OWT) dataset using the code from https://github.com/karpathy/nanoGPT, and LLAMA-2 130M for 20k and 350M for 60k steps on the C4 dataset following [1].  For these experiments, we adhered to the exact settings provided in the corresponding papers or GitHub repositories to ensure a fair comparison. The results are reported in Table 2 and Figure 1 of the rebuttal PDF. Similar to the vision tasks in the manuscript, our AdamW+4-bit Shampoo consistently outperformed AdamW and AdamW+4-bit Shampoo (naive) in terms of performance, and AdamW+32-bit Shampoo in terms of memory usage.  \n\nAdditionally, in the manuscript, we have evaluated six models (including VGG, ResNet, ViT, and Swin) on three datasets: CIFAR100, Tiny-ImageNet, and the large-scale ImageNet dataset. These diverse experiments sufficiently demonstrate the superiority of our 4-bit Shampoo over the vanilla 32-bit version.\n\nFinally, our algorithm design does not rely on any specific properties of the tasks, ensuring that the performance improvements are general and transferable. So we believe these improvements should be consistently achievable on other tasks as well. This is validated by the results on both the computer vision tasks in the manuscript and the natural language processing tasks presented in the rebuttal PDF.\n\n\n**2) Add a row for 8-bit A in Table 1.**\nPer your suggestion, we provide the quantization errors of 8-bit quantization schemes in Table 1 of the rebuttal PDF. Comparing Table 1 in the rebuttal PDF with Table 1 in the manuscript, we can see that 4-bit quantization of eigenvector matrix $U$ has smaller quantization errors than 8-bit quantization of the preconditioner $A$. For a clearer comparison, we extract the quantization errors using Linear-2 quantization from the two tables. The results are given in the following table, where $A = A_1$ is derived from the real world as described in Subsection 3.1. We will include these results into the revision. \n\n| Bits | Quantized Matrices | NRE  | AE ($^\\circ$) |\n| :--: | :----------------: | :--: | :-----------: |\n| 4 | $A$ | 0.6243 | 17.293 |\n| 4 | $U$ | 0.0543 | 3.1066 |\n| 8 | $A$ | 0.2164 | 7.9751 |\n| 8 | $U$ | 0.0037 | 0.2121 |\n\n**3) Add a description of DT quantization.**\nThanks. Dynamic tree (DT) quantization for $b$-bit quantization maps $\\\\{0, 1, \\dots, 2^b - 1\\\\}$ to $\\\\{0, 1\\\\}\\cup G$, where $G$ is a set of numbers with the following properties: the number in $G$ looks like $\\pm q_k \\times 10^{-E}$, where\n\na) $b = 2 + E + F$, where $E, F$ are integers;\n\nb) $q_k = (p_k + p_{k+1}) / 2$, where $k \\in \\\\{0, \\dots, 2^F-1\\\\}$;\n\nc) $ p_j = 0.9j / 2^F + 0.1$, where $j \\in \\\\{0, \\dots, 2^F\\\\}$. \n\nWe will add a paragraph in Appendix C for describing quantization maps mentioned in the manuscript.\n\n**4) Add a graph from $x \\in [-1,1]$ to $Q(x)$ for each quantization scheme.**\nPer your suggestion, we have drawn the graph from $x \\in [-1,1]$ to $Q(x)$ for each quantization scheme, which is similar to Figure 6 in [2]. Unfortunately, the space constraint prevents its inclusion in the one-page rebuttal PDF. We will add this graph in Appendix C.\n\n**5) Why separate $T_1$ and $T_2$ in Algorithm 3.**\nThanks. We follow previous second-order optimizers [3, 4], and set different values for $T_1$ and $T_2$ ($T_1$=200, $T_2$=2000 in [4]) for more efficient training. This is because the inverse root is much more expensive than preconditioner update, and thus is updated more lazily. Here we keep the default values of $T_1$ and $T_2$ for both 4- and 32-bit Shampoo for a fair comparison.\n\n**6) English improvement.**\nThank you for your meticulous proofreading. We will make every effort to polish our manuscript and also seek assistance from native English speakers.\n\n\n\n[1] Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. Galore: Memory-efficient llm training by gradient low-rank projection. ICML, 2024.\n\n[2] Tim Dettmers, Mike Lewis, Sam Shleifer, Luke Zettlemoyer. 8-bit optimizers via block-wise quantization. In ICLR, 2022.\n\n[3] Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer. Scalable second order optimization for deep learning. arXiv preprint arXiv:2002.09018, 2020.\n\n[4] Hongwei Yong, Ying Sun, and Lei Zhang. A general regret bound of preconditioned gradient method for DNN training. In CVPR, 2023.", "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": "Yztiy3veOf", "reviewer": "Reviewer_RMP9", "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.", "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. \nIn terms of timing the SoTA right now is nAdamW so Shampoo should be compared to that for wall clock times.  \n\nIdeally one would use the very recent schedule free optimization as a benchmark.\n\nhttps://github.com/facebookresearch/schedule_free\n\n If it is possible to add these as a benchmarks with timings it would be more convincing. \n\nFurthermore, 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": "As mentioned above the paper does not consider LLMs which can often have a large impact on memory requirement vs a standard vision models. \n\nThe paper does not consider more SoTA optimizers like nAdamW or LARS, and optimization methods like EMA and lookahead.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The method seems to regain performance vs just naively quantizeing the preconditioner. This makes Shampoo require about the same about of memory as Adam (at least for the vision experiments in the paper).\n\nThey presented the theoretical reason why quantizing the eigen matrix is better and basically showed that the quantization is a form of noise and one can just basically find the closest orthonormalization to recover the proper eigen directions via  Björck orthonormalization.", "weaknesses": "Could benefit from extensions to LLMs:\n\nI 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. \n\nLearning rate and Learning rate Schedule:\n\nFor 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.\n\nFor the ViT Imagenet 1-k the initial lr should likely be better tuned.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the insightful and positive comments. In the following, we provide our point-by-point response and hope our response helps address your concerns. We also look forward to the subsequent discussion which further helps to solve the current issues.\n\n**1) Add experiments on LLMs \\& 124M GPT2.**\nThanks. Due to limited computing resources and the short rebuttal period of 7 days, we trained medium-sized language models, including 124M GPT-2 for 20k steps on the OpenWebText (OWT) dataset using code from nanoGPT (https://github.com/karpathy/nanoGPT), and LLAMA-2 130M for 20k and 350M for 60k steps on the C4 dataset following [1]. For these experiments, we adhered to the exact settings provided in the corresponding papers or GitHub repositories to ensure a fair comparison. The results are reported in Table 2 and Figure 1 of the rebuttal PDF. As with the vision tasks in the manuscript, our AdamW+4-bit Shampoo consistently outperformed AdamW and AdamW+4-bit Shampoo (naive) in terms of performance, and AdamW+32-bit Shampoo in terms of memory usage. Since our algorithm design is not task-specific, we believe its improvements should be consistently achievable on other tasks as well which is also validated by the results on both the computer vision tasks in the manuscript and the natural language processing tasks here.\n\nFor LLMs, we have compared our quantization with the naive quantization approach in Table 2 and Figure 1 of the rebuttal PDF. Notably, our 4-bit Shampoo with our quantization outperforms the naive 4-bit Shampoo with naive quantization in terms of validation loss with almost the same training times and memory usage. This demonstrates that the performance gain of our quantization is task-independent. \n\n**2) More proper learning rate decay for ResNet.**\nThanks. Due to time limitation, we use cosine learning rate decay (initial lr=0.1) to train ResNet34 on CIFAR-100. Experimental results can be found in the table below. By comparison, SGDM+Shampoo still converges faster than SGDM, and have slightly better test performance.\n\n| Epochs |        Optimizer         | TA (%) | WCT (min) |\n| :----: | :----------------------: | :----: | :-------: |\n|  200   |           SGDM           | 79.67  |   116.0   |\n|  300   |           SGDM           | 79.83  |   172.7   |\n|  200   |   SGDM+32-bit Shampoo    | 80.39  |   152.7   |\n|  200   | SGDM+4-bit Shampoo (our) | 80.22  |   161.7   |\n\n**3) Initial learning rate for ViT ImageNet 1-k.**\nThanks. We use the default initial learning rate provided in [2] to train ViT-Base. Although there may be a better learning rate, the performance reported in our paper is relatively high and reasonable. For ViT-Base/32 on ImageNet 1-k, Table 5 in [3] reports 73.38% accuracy for training 300 epochs, while ours for 150 epochs is 72.87% accuracy. \n\n**4) Comparison with NadamW and schedule free optimization.**\nPer your suggestion, we provide the results of training Swin-Tiny on CIFAR-100 in the table below. We run AdamW, NadamW, and AdamWScheduleFree for 150 epochs and AdamW+Shampoo for 100 epochs. For NadamW, we use its implementation in Timm library, using the same hyperparameters as AdamW. For schedule free optimization, we use the code from [4] to train Swin-Tiny, and set lr=0.0025 (default), weight_decay=0.05, warmup_steps=10000. \n\n|         Optimizer         | TA (%) | WCT (min) |\n| :-----------------------: | :----: | :-------: |\n|           AdamW           | 76.69  |   318.6   |\n|          NadamW           | 77.11  |   342.4   |\n|     AdamWScheduleFree     | 76.58  |   321.9   |\n| AdamW+4-bit Shampoo (our) | 78.63  |   273.3   |\n\nOne can see that though improving AdamW, these methods are still worse than our AdamW+4-bit Shampoo.\n\n\n\n[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.\n\n[2] Pan Zhou, Xingyu Xie, and Shuicheng Yan. Win: Weight-decay-integrated Nesterov acceleration for adaptive gradient algorithms. In ICLR, 2023.\n\n[3] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In ICLR, 2021.\n\n[4] The Road Less Scheduled. arXiv preprint arXiv:2405.15682, 2024.", "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:
We thank all the reviewers for their valuable comments and suggestions. In the attached PDF, we have included additional results to address the reviewers' concerns, specifically in Table 1, Table 2, and Figure 1. We also give Lemma 1 to prove the convergence of our 4-bit Shampoo. We will incorporate these results and lemma in the appendix of the revised manuscript.

**Table 1:**
As mentioned in the conclusion section of the manuscript, preconditioners in Shampoo are symmetric matrices and thus can be stored as upper triangular matrices, saving almost half of the memory usage. However, the eigenvector matrix of a preconditioner is not symmetric, resulting in an 8-bit preconditioner occupying the same memory as its 4-bit eigenvector matrix. Comparing Table 1 in the attached PDF with Table 1 in the manuscript, one can observe that the 4-bit quantization of the eigenvector matrix $U$ has smaller quantization errors than the 8-bit quantization of the preconditioner $A$. For a clearer comparison, we extract the quantization errors using Linear-2 quantization from the two tables. The results are given in the following table, where $A = A_1$ is derived from real-world data as described in Subsection 3.1 of the manuscript.

| Bits | Quantized Matrices | NRE  | AE ($^\circ$) |
| :--: | :----------------: | :--: | :-----------: |
| 4 | $A$ | 0.6243 | 17.293 |
| 4 | $U$ | 0.0543 | 3.1066 |
| 8 | $A$ | 0.2164 | 7.9751 |
| 8 | $U$ | 0.0037 | 0.2121 |

**Table 2 and Figure 1:**
In the manuscript, we provided diverse experiments on computer vision tasks to demonstrate that the performance improvements of our approach are general and transferable. Due to limited computing resources and the short rebuttal period of 7 days, we trained medium-sized language models to address the reviewers' concerns. Specifically, we trained a 124M GPT-2 for 20k steps on the OpenWebText (OWT) dataset using code from nanoGPT (https://github.com/karpathy/nanoGPT), and LLAMA-2 130M for 20k and 60k steps on the C4 dataset following [1]. Unless otherwise noted, the maximum order of a preconditioner for training GPT-2 is 1200, and for training LLAMA-2 is 10000. As with the vision tasks in the manuscript, our AdamW+4-bit Shampoo consistently outperformed AdamW and AdamW+4-bit Shampoo (naive) in terms of performance, and AdamW+32-bit Shampoo in terms of memory usage. Since our algorithm design is not task-specific, we believe its improvements should be consistently achievable on other tasks as well. 

**Lemma 1:**
We find that the proof techniques of Theorem 1 in [2] can be adapted to demonstrate the convergence of our optimizer. Specifically, we can generalize Lemma 2 used in proving Theorem 1 in [2] and then use this extended lemma to prove the convergence of quantized Shampoo. Our extension of Lemma 2 in [2] is the Lemma 1  below. Providing a complete and detailed convergence analysis within a short 7-day rebuttal period is challenging, but we will make every effort to include this analysis in the revised manuscript.

**Lemma 1.**  Let $\\{X_t\\}$ be a sequence of symmetric matrices, and $A_t=\sum_{s=1}^tX_s$, where $t=1, \dots, T$. Suppose we have two sequences of symmetric matrices $\\{Y_t\\}$, $\\{ Z_t \\}$, and a sequence of real numbers $\\{\rho_t\\}$ satisfying
$$
Y_t = Z_{t-1}+X_t, \quad \rho_t=\rho_{t-1}+\|\| Y_t - Z_t\|\|, \quad Z_0=0, \rho_0=0.
$$
Define $B_t=\rho_tI+Z_t$, where $I$ denotes the identity matrix. Then for $t=1, \dots, T$, we have
$$
B_t \succeq B_{t-1} + X_t, \quad A_t \preceq B_t \preceq 2\rho_tI + A_t.
$$
*Proof.*  Note that for any symmetric matrix $S$, it holds that $\|\|S\|\|I\succeq S$. Then we have
$$
(\rho_t-\rho_{t-1})I + Z_t = \|\| Y_t - Z_t\|\|I + Z_t \succeq Y_t.
$$
Adding $\rho_{t-1}I$ on both sides, we get
$$
B_t = \rho_tI + Z_t \succeq \rho_{t-1}I + Y_t = \rho_{t-1}I + Z_{t-1}+X_t = B_{t-1} + X_t.
$$
Hence
$$
B_t = \sum_{s=1}^{t}(B_s-B_{s-1}) \succeq \sum_{s=1}^{t}X_s = A_t.
$$
On the other hand, we have
$$
Z_t \preceq \|\| Z_t - Y_t\|\|I + Y_t = (\rho_t-\rho_{t-1})I + Y_t.
$$
Adding $\rho_t I$ on both sides, we get
$$
B_t  = \rho_t I + Z_t \preceq (2\rho_t-\rho_{t-1})I + Y_t
$$
$$
		 = 2(\rho_t-\rho_{t-1})I + \rho_{t-1}I + Z_{t-1}+X_t
$$
$$
		 = B_{t-1} + 2(\rho_t-\rho_{t-1})I + X_t.
$$
Hence
$$
B_t = \sum_{s=1}^{t}(B_s-B_{s-1}) \preceq \sum_{s=1}^{t}2(\rho_s-\rho_{s-1})I + \sum_{s=1}^{t}X_s = 2\rho_tI+ A_t.
$$
The proof is completed.

[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.

[2] Jui-Nan Yen, Sai Surya Duvvuri, Inderjit S. Dhillon, and Cho-Jui Hsieh. Block low-rank preconditioner with shared basis for stochastic optimization. NeurIPS, 2023.

Meta-review: 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.
Decision: Accept (poster)