Title: 4-bit Shampoo for Memory-Efficient Network Training
Venue: NeurIPS 2024
Accepted: yes
Keywords: memory efficient

second-order optimizer

Shampoo

quantization
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-Efficient Network Training

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

1 Introduction

Deep neural networks (DNNs) have achieved great success in numerous fields, e.g., computer vision [20], natural language processing [38], and speech recognition [16]. A significant part of such success is attributed to first-order optimizers such as stochastic gradient descent with momentum (SGDM) [31] and AdamW [29]. Second-order optimizers, including K-FAC [30], Shampoo [18], AdaBK [41], CASPR [13], and Sophia [27], show great convergence properties, but often involve noticeable computation and memory costs. Anil et al. [2] provided several practical techniques for second-order optimizers to achieve substantial wall-clock time improvements over traditional first-order optimizers. The fast convergence property of second-order optimizers benefits from preconditioning the gradient with a matrix known as a preconditioner. The optimizer states for constructing the preconditioner and its inverse root can speed up optimization compared to first-order optimizers, but consume memory that could be used for model parameters, limiting the maximum model size trained within a given memory budget. With the increase in model size, the memory utilized by optimizer states can become a predominant factor in memory usage. This is the primary obstacle hindering the widespread use of second-order optimizers in the era of large models.

There are two main attempts to reduce memory consumed by optimizer states. Factorization uses lowrank approximation to optimizer states. This strategy has been applied to first-order optimizers [35, 3]

0

50

100

150

200

250

300

Wall-clock Time (min)

50

55

60

65

70

75

80

Test Accuracy (%)

AdamW

AdamW+32-bit Shampoo

AdamW+4-bit Shampoo (naive)

AdamW+4-bit Shampoo (our)

800

1000

1200

1400

1600

1800

2000

2200

1465.8

2036.0

1543.9

1543.9

AdamW

AdamW+32-bit Shampoo

AdamW+4-bit Shampoo (naive)

AdamW+4-bit Shampoo (our)

(a) Swin-Tiny on CIFAR-100

0

500

1000

1500

2000

Wall-clock Time (min)

30

40

50

60

70

Test Accuracy (%)

AdamW

AdamW+32-bit Shampoo

AdamW+4-bit Shampoo (naive)

AdamW+4-bit Shampoo (our)

9000

9500

10000

10500

11000

11500

12000

12500

10600

12134

10803

10804

AdamW

AdamW+32-bit Shampoo

AdamW+4-bit Shampoo (naive)

AdamW+4-bit Shampoo (our)

(b) ViT-Base/32 on ImageNet-1k

and second-order optimizers [14, 40]. In a comparable but distinct line of work, quantization utilizes low-bit to compress 32-bit optimizer states. Quantization is attractive due to its simplicity and wide applicability, which has been applied to first-order optimizers [8, 26]. Applying quantization to second-order optimizers poses a greater challenge, as first-order optimizers' states are elementwise, whereas second-order optimizers rely on matrix operations. To our knowledge, it has not been attempted before.

Contributions: In this paper, we present the first second-order optimizers with 4-bit optimizer states by taking Shampoo [18] as an example, while preserving the performance achieved with 32-bit optimizer states. While our focus is on Shampoo, we believe that our approach could also be applied to other second-order optimizers (see Table 4). Our main contributions are highlighted below.

Firstly, to maintain 32-bit performance, we propose quantizing the eigenvector matrix of a preconditioner in 4-bit Shampoo, rather than the preconditioner itself. The reason is that the small singular values of the preconditioner matter. Directly quantizing the preconditioner via block-wise quantization [8] at 4-bit precision can significantly alter the small singular values, leading to a drastic change in its inverse 4-th root and thus harming 4-bit Shampoo's performance. Quantizing the eigenvector matrix can help alleviate this issue, which is supported by experimental validation and theoretical insight. Additionally, with the eigenvector matrix, computing the inverse 4-th root is straightforward, ensuring that quantizing the eigenvector matrix does not lead to a rise in the total wall-clock time compared to quantizing the preconditioner (see Figure 1).

Secondly, we present two techniques for enhancing performance. As the eigenvector matrix of a preconditioner is orthogonal, we apply Björck orthonormalization [4] to rectify the orthogonality of the quantized eigenvector matrix, leading to improved approximation of preconditioner's eigenvector matrix and facilitating computation of its inverse 4-th root. Additionally, we observe that linear square quantization outperforms dynamic tree quantization [7] marginally when quantizing second-order optimizer states. The superiority of our developed 4-bit Shampoo is demonstrated in Figure 1.

Finally, we evaluate our 4-bit Shampoo on different image classification and natural language modeling tasks using convolutional neural network (CNN) and transformer architectures. Across all these benchmarks, our 4-bit Shampoo achieves similarly fast convergence comparable to its 32-bit counterpart, with no significant increase in losses for the trained models. Our 4-bit Shampoo uses less memory than its 32-bit counterpart, allowing for training of larger models with given resources.

2 Preliminaries

In this section, we present Shampoo and its implementation in our experiments. We also discuss quantization-based compression methods in a general formulation.

Notations. We use a non-bold letter like a or A to denote a scalar, a boldfaced lower-case letter like a to denote a vector, and a boldfaced upper-case letter such as A to denote a matrix. u =[ u i ] T means that the i -th element of column vector u is u i and U =[ u i ] means the i -th column vector of matrix U is u i . Let A be a positive definite (PD) matrix and s ∈ R , we define A s = U Λ s U T , where U Λ U T is the Singular Value Decomposition (SVD) of A . tr( A ) represents the trace of a matrix A . The inner product of two matrices A and B is denoted as ⟨ A , B ⟩ =tr( A T B ) . The Frobenius norm of a matrix A is ∥ A ∥ F = √ ⟨ A , A ⟩ . A ⊙ B means the elementwise matrix product (Hadamard product).

Diag( a ) is a diagonal matrix with diagonal vector a , while diag( A ) means the diagonal vector of matrix A .

2.1 Shampoo for Matrices

The update rule of Shampoo in the matrix case combined with a first-order optimizer F is

where W t is the model parameters in matrix form, L t and R t are called preconditioners, s t is the optimizer state of F , and G t is the gradient at W t -1 . Note that L t , R t , L -1 / 4 t , and R -1 / 4 t are PD matrices. The penultimate step in (1) is the grafting trick [1], which enables Shampoo to roughly apply the well-tuned learning rate schedule of F . The optimization variable W t does not represent all model parameters. It denotes a tensor of the model [18] or one block of a tensor [2]. In practice, we adopt an efficient and effective implementation of Shampoo for training DNNs following [2, 41] as described in Algorithm 4. In order to achieve efficient training, L t , R t , L -1 / 4 t , and R -1 / 4 t are computed once every few hundred iterations. In this case, besides L t and R t , their inverse 4-th roots should also be stored in memory, as computing them is computationally expensive. So training large models with Shampoo can be memory-intensive, consuming a significant amount of memory.

2.2 Quantization-based Compression Methods

Quantizing updated optimizer states using a quantizer and then dequantizing them with a dequantizer prior to use is an effective method for conserving memory. We focus exclusively on vectors, as tensors can be reshaped into vectors.

Quantization. According to the idea in [8, 26], a b -bit quantizer Q for p -dimensional real vectors is a mapping given by

where N is a normalization operator on R p , I is an elementwise function mapping any real number to an element of T b = { 0 , 1 , . . . , 2 b -1 } , and M is a maximum operator on R p . For any x ∈ R p , N and M satisfy N ( x ) ⊙M ( x )= x .

A normalization operator N for p -dimensional vectors is a transformation on R p . It scales each element of a vector x ∈ R p into [ -1 , 1] . A block-wise normalization operator for a p -dimensional vector x = [ x 1 , x 2 , . . . , x p ] T is defined as

where N ( x ) i is the i -th element of N ( x ) , and X i is a set satisfying i ∈ X i ⊂ { 1 , . . . , p } . Usually, X i should also satisfy X i = X j or X i ∩ X j = ∅ for i, j ∈ { 1 , . . . , p } . In this case, for any x ∈ R p , the number of different elements in M ( x ) is equal to the number of elements in set { X i | i = 1 , . . . , p } . Meanwhile, the number of the elements in X i for any i should be as close as possible to a value called block size.

The mapping I for x ∈ R in a b -bit quantizer Q is defined as

where R named quantization mapping is an elementwise function that maps any element in T b into [ -1 , 1] , and | · | is the absolute operator for a scalar. There are three typical quantization mappings: linear quantization, dynamic quantization, and quantile quantization. Their specifications and visualizations can be found in [8].

Dequantization. Given a b -bit quantizer Q =( I ◦ N , M ) for a p -dimensional real vector x ∈ R p , the corresponding dequantizer D is a mapping defined as

3 Methodology

In this section, we describe the design of our quantization-based compression method to realize 4-bit Shampoo with fast and high precision quantization. Let Q =( I ◦ N , M ) be a quantizer and D be its corresponding dequantizer as described in Subsection 2.2.

3.1 Quantizing the Eigenvector Matrices

A naive approach to realize 4-bit Shampoo is applying the compression methods proposed in [8, 26] to L t , R t , L -1 / 4 t , and R -1 / 4 t in Shampoo (see (1)). A slightly improved approach is to quantize the four PD matrices excluding their diagonal elements, which are typically much larger than their nondiagonal counterparts due to the non-negativity of the elements in diag( G t G T t ) and diag( G T t G t ) .

However, the naive approach can cause large quantization errors at 4-bit precision. This is because the quantization errors (or called perturbations) of quantizing L t and R t will transfer to L -1 / 4 t and R -1 / 4 t . To verify this, we first introduce two criteria to evaluate the quantization errors of matrices. We do not use the elementwise criterion in [8]. Let A denote a 32-bit matrix, g represent a transformation (can formed by quantization), and f stand for a mapping, e.g., f ( A )= A -1 / 4 . Then we define the normwise relative error (NRE) and angle error (AE) in f of g at A as

We choose two PD matrices of order 1200. The first one A 1 is derived from the real world. It is a preconditioner in 32-bit Shampoo combined with AdamW for training a Swin-Tiny model. The second one A 2 = U Λ U T is synthetic, constructed from a random orthogonal matrix U and a diagonal matrix Λ with only two distinct diagonal values. Table 1 shows the quantization errors in f ( A ) = A -1 / 4 of the naive approach at these two matrices, which are remarkably high. More analyses are given in Appendix D. The key point is that the singular values of A i ( i =1 , 2) follow a specific distribution (see Figure 2). In this scenario, a slight perturbation of A i will significantly alter its small singular values, resulting in a drastic change to A -1 / 4 i .

To address this issue, we propose quantizing the eigenvector matrix of a preconditioner in Shampoo, rather than the preconditioner itself. Namely, a preconditioner A is a PD matrix, and its SVD is U Λ U T , where U represents the eigenvector matrix and Λ denotes the singular value matrix. Given that Λ is a diagonal matrix, we can focus on quantizing U using Q while leaving Λ unchanged. From Table 1, one can observe that quantizing U can significantly reduce the quantization errors. We will theoretically discuss the advantages of quantizing U compared to quantizing A in Section 4. In practice, the randomized SVD method [19] is adopted to compute the SVD of A efficiently, as shown in [40]. We want to highlight that quantizing the original L t and R t in Shampoo involves significant computational burdens to compute their inverse 4-th roots L -1 / 4 t and R -1 / 4 t , whereas quantizing the eigenvector matrices of L t and R t allows for rapid inverse root calculation. So the computational time required for both approaches is comparable (see Figure 1).

0

200

400

600

800 1000 1200

Index

11

10

9

8

7

6

5

Singular Value

real

quan

(a) Real-world

0

200

400

600

800 1000 1200

Index

5

4

3

2

1

0

Singular Value

real

quan

(b) Synthetic

0

1

2

3

4

5

Number of Iterations

t

2

6

5

4

3

2

1

0

1

2

Mean Error

s

=

1

s

=

1/4

s

=

1/10

s

=

1/20

3.2 Rectifying the Orthogonality of Eigenvector Matrices

Let A be a PD matrix with SVD U Λ U T . Note that the eigenvector matrix U is orthogonal, whereas V = D ( Q ( U )) may not be. To further mitigate the quantization errors mentioned in Subsection 3.1, we propose employing Björck orthonormalization [4] to orthogonalize V . Particularly, given V 0 = V , we iterate

for t 1 ≥ 1 times and take V t 1 as the rectified result. Equation (2) can also be interpreted as the gradient descent of problem min V ∥ V T V -I ∥ 2 F using a step size of 0.5, where I denotes the identity matrix. We empirically find that only one iteration (i.e., t 1 =1 ) is enough. Table 1 illustrates the benefit of rectifying V into V 1 .

The update frequencies for the preconditioners and their inverse 4-th roots differ (see Algorithm 3). Given V and Λ , we also require orthogonal rectification to compute A s rapidly for any s ∈ R . The reason is as follows. It is easy to compute A s = U Λ s U T by definition. However, U Λ s U T can be very sensitive to the orthogonality of U for s < 0 , making V Λ s V T largely deviate from ( V Λ V T ) s ≈ A s . Similarly, we can approximate A s by V t 2 Λ s V T t 2 , where V t 2 is generated by (2). Figure 3 illustrates the elementwise mean errors between ( V t 2 Λ s V T t 2 ) -1 / s ( V t 2 Λ V T t 2 ) and I for various s and t 2 , where A is the real-world matrix used in Table 1. Based on the observation from Figure 3, we set t 2 =4 in our experiments.

3.3 Selecting the Quantizer

The quantizer Q is defined by the normalization operator N and mapping R , and N is determined by X i . Since an eigenvector has a unit length, the elements in X i should belong to the same column of an eigenvector matrix, i.e., they are from the same eigenvector. Instead of employing dynamic tree (DT) quantization as mapping R , we recommend utilizing linear square (Linear-2) quantization as R , particularly when b =4 . Linear-2 quantization is defined as

where j ∈ T b = { 0 , 1 , . . . , 2 b -1 } . As shown in Table 1, Linear-2 quantization has lower quantization errors compared to DT quantization at 4-bit precision.

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 a PD matrix and its SVD is U Λ U T . We can compress A into a pair ( λ , U ) = (diag( Λ ) , Q ( U )) and decompress it into ( Λ , V ) = (Diag( λ ) , D ( U )) . Algorithm 1 (Preconditioner Update, PU) shows the update rule of A . Similarly, we compress ̂ A ≈ A -1 / 4 into a pair ( a , A ) = (diag( ̂ A ) , Q ( ̂ A -Diag( a ))) and decompress it into Diag( a ) + D ( A ) . Algorithm 2

(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 in Algorithm 3. Note that we omit some input parameters of PU and PIRU because they can be found in Algorithm 3 in the same form.

Algorithm 3 Practical 4-bit Shampoo

Input: W 0 ∈ R m × n , L 0 = ϵ I m , R 0 = ϵ I n , ̂ L 0 = I m , ̂ R 0 = I n , β ∈ (0 , 1) , t 1 , t 2 , update interval T 1 , update interval T 2 , total number of steps T , first-order optimizer F , first-order optimizer state s 0 = 0 , 4-bit quantizer Q and its corresponding dequantizer D .

Output: final parameter W T .

2:

l

0

= diag(

̂

L

0

)

,

L

0

=

3: for t = 1 , 2 , . . . , T do

4: Receive loss function L t : R m × n ↦→ R and compute gradient G t = ∇L t ( W t )

5: if t % T 1 ≡ 0 then

6:

t,L

t,L

λ

U

,

=PU(

7: else

8:

λ

t,L

9:

if

t

%

,

U

≡

t,L

0

=

λ

then

t

λ

-

t

-

1

,L

1

,

,L

U

,

t

U

-

t

-

1

,L

1

;

,

G

λ

t

G

t,R

,

t

T

U

);

λ

t,R

t,R

=

,

λ

U

t

-

t,R

1

,R

=PU(

,

U

t

-

λ

t

1

,R

2 10: l t , L t = PIRU( λ t,L , U t,L ); r t , R t = PIRU( λ t,R , U t,R )

11: else

13: ̂ L t = Diag( l t ) + D ( L t ); ̂ R t = Diag( r t ) + D ( R t )

14: ̂ G t = ̂ L t G t ̂ R t ; ˜ G t = ̂ G t ( ∥ G t ∥ F / ∥ ̂ G t ∥ F )

15: W t , s t = F ( W t -1 , s t -1 , ˜ G t )

4 Theoretical Analysis

In this section, we analyze why quantizing the eigenvector matrix of a preconditioner in Shampoo is better than quantizing the preconditioner itself under a certain singular value distribution. Furthermore, we consider quantization as a perturbation and prove the convergence of the perturbed Shampoo (Algorithm 6) in Appendix E. The following lemma reveals some good properties of perturbing the eigenvector matrix of a PD matrix.

Lemma 1. Let A be a PD matrix whose SVD is U Λ U T , where U =[ u i ] is an orthogonal matrix and Λ =diag([ λ i ] T ) is a diagonal matrix. Given a perturbation ∆ U =[∆ u i ] and s ∈ R , we define B :=( U Λ U T ) s and ∆ B :=(( U +∆ U ) Λ ( U +∆ U ) T ) s -B .

If U +∆ U is orthogonal and there exists α ∈ R such that ∥ ∆ u i ∥ 2 ≤ α , then

If U +∆ U is orthogonal and there exists β ∈ R such that ⟨ u i , u i +∆ u i ⟩ ≥ 1 -β ≥ 0 , then

T

Q

(

0

);

r

0

= diag(

R

̂

0

)

,

R

0

=

Q

(

0

)

,L

-

1

,R

,

U

t

-

1

,R

,

G

T

t

G

t

)

From Lemma 1, it is evident that the normwise relative error and angle error in f ( A ) = A s of perturbing U at A = U Λ U T are independent of Λ and s . Moreover, these errors are well-bounded under some mild conditions. Empirically, for 4-bit quantization, α = 0 . 1 and β = 0 . 005 roughly meet the conditions of Lemma 1, leading to ∥ ∆ B ∥ F ∥ B ∥ F ≤ 0 . 2 and ⟨ B , B +∆ B ⟩ ∥ B ∥ F ∥ B +∆ B ∥ F ≥ 0 . 99 .

It is very complicated to generally analyze the perturbation in f ( A ) = A s of perturbing A . Thus, we focus on perturbing the singular values of A . For simplicity, we assume that both A and A +∆ A have only two distinct singular values, where ∆ A is a perturbation of A . The following lemma gives the perturbation in A s of perturbing the smaller singular value of A .

Lemma 2. Let A be a PD matrix of order m + n whose SVD is U Λ U T , where m,n ∈ N + , n = lm , U = [ u i ] is an orthogonal matrix and Λ = diag([ λ i ] T ) is a diagonal matrix. Assume that Λ = diag([ cλ 1 T m × 1 , λ 1 T n × 1 ] T ) , c ≥ 1 , and λ > 0 . Given a perturbation ∆ Λ = diag([ 0 T m × 1 , ∆ λ T n × 1 ] T ) and s ∈ R , we define B :=( U Λ U T ) s and ∆ B :=( U ( Λ +∆ Λ ) U T ) s -B .

Moreover, h 1 ( s, l ) decreases monotonically with s over ( -∞ , 0) and increases monotonically with l over (0 , + ∞ ) .

If ∆ λ n × 1 = ( tc -1) λ 1 n × 1 where t > 0 , then

Moreover, h 2 ( l ) decreases monotonically with l over (0 , ( c/t ) s ] and increases monotonically with l over (( c/t ) s , + ∞ ) .

(3) If ∆ λ n × 1 = ( tc -1) λ 1 n × 1 where k = tc > 0 and l = ( c/t ) s , then

Let us make some comments on the above lemma. First, from Lemma 2(1) we have h 1 (1 , l ) = ∥ ∆ A ∥ F ∥ A ∥ F = √ l | k -1 | √ c 2 + l . If k ≥ 1 , ∥ ∆ A ∥ F ∥ A ∥ F = ∥ ∆ Λ ∥ F ∥ Λ ∥ F is bounded by k c √ l = t √ l . Second, if k = tc ≥ 1 and s < 0 , one can deduce h 2 ( l ) ≥ √ lt 2 s / (1 + lt 2 s ) from Lemma 2(2), which indicates that a small lt 2 s is needed to achieve small h 2 ( l ) . We can set t = 0 . 02 to simulate 4-bit quantization. Based on Lemma 1 and Lemma 2(3), we have the following proposition.

Proposition 1. Let A be a PD matrix of order m + n whose SVD is U Λ U T , where m,n ∈ N + , n = lm , U =[ u i ] is an orthogonal matrix, Λ =diag([ cλ 1 T m × 1 , λ 1 T n × 1 ] T ) , c ≥ 1000 , and λ > 0 . Given ∆ U =[∆ u i ] , ∆ Λ =diag([ 0 T m × 1 , ∆ λ T n × 1 ] T ) , and s ≤-0 . 25 , we define B := ( U Λ U T ) s , B 1 := (( U + ∆ U ) Λ ( U +∆ U ) T ) s , and B 2 := ( U ( Λ +∆ Λ ) U T ) s . If U + ∆ U is orthogonal, ∥ ∆ u i ∥ 2 ≤ 0 . 1 , ⟨ u i , ∆ u i ⟩≥-0 . 005 , ∆ λ n × 1 =(0 . 02 c -1) λ 1 n × 1 , and l =( c/ 0 . 02) s , then

Proposition 1 requires very strong assumptions. Nevertheless, it provides insight into why quantizing A can result in a greater normwise relative error and angle error in A s , compared to quantizing U . Complete proofs of Lemma 1, Lemma 2, and Proposition 1 can be found in Appendix F.

5 Experiments

In this section, we compare our 4-bit Shampoo combined with SGDM or AdamW to their 32-bit counterparts, as well as the first-order optimizers on various image classification tasks. See more experimental results on image classification and natural language modeling tasks in Appendix H.

Models, datasets, and hyperparameters. We train VGG19 [36], ResNet34 [20], ViT-Small [10], and Swin-Tiny [28] on the CIFAR-100 [23] and Tiny-ImageNet [24] datasets with one RTX3060Ti GPU, and train ResNet50 and ViT-Base/32 on the ImageNet-1k dataset [34] with one A800 GPU.

0

25

50

75

100 125 150 175

Wall-clock Time (min)

50

55

60

65

70

75

80

Test Accuracy (%)

ResNet34 on CIFAR-100

SGDM

SGDM+32-bit Shampoo

SGDM+4-bit Shampoo

0

100 200 300 400 500 600 700

Wall-clock Time (min)

50

55

60

65

70

75

80

Test Accuracy (%)

ViT-Small on CIFAR-100

AdamW

AdamW+32-bit Shampoo

AdamW+4-bit Shampoo

0

500

1000

1500

2000

Wall-clock Time (min)

30

40

50

60

70

80

Test Accuracy (%)

ResNet50 on ImageNet-1k

SGDM

SGDM+32-bit Shampoo

SGDM+4-bit Shampoo

0

500

1000

1500

2000

Wall-clock Time (min)

30

40

50

60

70

Test Accuracy (%)

ViT-Base/32 on ImageNet-1k

AdamW

AdamW+32-bit Shampoo

AdamW+4-bit Shampoo

For hyperparameter settings, we mainly follow [41] to train CNNs and [25, 44] to train vision transformers. For all the tasks, we keep the common hyperparameters of optimizers the same values. See Appendix G for experimental details.

Main results. We show the performance, wall-clock time, and memory cost in Table 2. Firstorder optimizers run 1.2x to 1.5x epochs, resulting in longer wall-clock time, yet yielding lower test accuracies compared to second-order optimizers. In comparison to 32-bit Shampoo, our 4-bit Shampoo shows comparable test accuracies with differences ranging from -0.7% to 0.5%, increases in wall-clock time varying from -0.2% to 9.5%, and memory savings of 4.5% to 41%. Compared to the first-order optimizers, the memory costs of our 4-bit Shampoo only rise by 0.8% to 12.7%. This represents a significant advancement in the utilization of second-order optimizers. Following [26], we report the total peak GPU memory consumption rather than the optimizer's peak GPU memory consumption. Our main focus is on quantizing the states for constructing preconditioners and their inverse roots, which are approximately 7x smaller for 4-bit Shampoo compared to 32-bit Shampoo (see Appendix G). Figure 4 shows the test accuracy curves on the CIFAR-100 and ImageNet-1k

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. Weinvestigate the effectiveness of our proposed quantization techniques. Table 3 indicates that quantizing the eigenvector matrix of a preconditioner is crucial for b -bit ( b = 3 , 4 ) Shampoo to maintain 32-bit performance, and orthogonal rectification is highly beneficial for 3-bit Shampoo. As for quantization mapping, linear square (Linear-2) quantization is comparable to dynamic tree (DT) quantization. We further apply our 4-bit quantization techniques to K-FAC [30], AdaBK [41] and CASPR [13] and the results are shown in Table 4. We can see that the 4-bit optimizers match the performance of their 32-bit counterparts, and reduce memory by over 20%.

6 Related Work

Second-order optimizers. Different second-order optimizers apply different second-order information. Hessian-based optimizers [39, 27] use the Hessian matrix or its approximation. Fisher-based optimizers [30, 41] utilize the covariance matrix of the accumulated gradients or its approximation based on Kronecker product. Shampoo [18] and CASPR [13] approximate the full AdaGrad [12] preconditioner by a set of small preconditioning matrices.

Memory efficient optimizers based on factorization. Adafactor [35] employs the outer product of two vectors to approximate the second moment of Adam [22]. SM3 [3] considers approximating the second moment of Adam by its covers' statistics. [14] and [40] reduce memory cost of the preconditioner in a second-order optimizer with its low-rank approximation through truncated SVD.

Memory efficient optimizers based on quantization. Dettmers et al. [8] introduce block-wise dynamic quantization that enables the use of first-order optimizers with 8-bit states. Li et al. [26] push the optimizer states of Adam/AdamW to 4-bit.

7 Conclusions, Limitations, and Broader Impact

We propose 4-bit Shampoo, the first low-bit second-order optimizer, designed for memory-efficient training of DNNs. We find that quantizing the eigenvector matrix of the preconditioner is essential to minimize quantization errors in its inverse 4-th root at 4-bit precision, given its sensitivity to alterations in small singular values. We further introduce orthogonal rectification and linear square quantization mapping to improve performance. 4-bit Shampoo achieves lossless performance to 32-bit counterpart in training different DNNs on various tasks.

Limitations. Preconditioners in Shampoo are symmetric matrices and can be stored as upper triangular matrices, saving almost half of the memory usage. However, the eigenvector matrix of a preconditioner is not symmetric, causing an 8-bit preconditioner to occupy the same memory as its 4-bit eigenvector matrix. Notably, a comparison of Table 1 and Table 7 in Appendix D 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 in 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-order optimizers. This could open up new research possibilities that were previously unattainable due to 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 1 grants (project ID: 23-SISSMU-028 and 23-SIS-SMU-070).

References

Naman Agarwal, Rohan Anil, Elad Hazan, Tomer Koren, and Cyril Zhang. Disentangling adaptive gradient methods from learning rates. arXiv preprint arXiv:2002.11803 , 2020.

Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer. Scalable second order optimization for deep learning. arXiv preprint arXiv:2002.09018 , 2020.

Rohan Anil, Vineet Gupta, Tomer Koren, and Yoram Singer. Memory efficient adaptive optimization. Advances in Neural Information Processing Systems , 32, 2019.

Å. Björck and C. Bowie. An iterative algorithm for computing the best estimate of an orthogonal matrix. SIAM Journal on Numerical Analysis , 8(2):358-364, 1971.

R.L. Burden, J.D. Faires, and A.M. Burden. Numerical Analysis . Cengage Learning, 2015.

Aaron Defazio, Xingyu Yang, Harsh Mehta, Konstantin Mishchenko, Ahmed Khaled, and Ashok Cutkosky. The road less scheduled. arXiv preprint arXiv:2405.15682 , 2024.

Tim Dettmers. 8-bit approximations for parallelism in deep learning. In Proceedings of the International Conference on Learning Representations , 2016.

Tim Dettmers, Mike Lewis, Sam Shleifer, and Luke Zettlemoyer. 8-bit optimizers via blockwise quantization. In Proceedings of the International Conference on Learning Representations , 2022.

Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, and Luke Zettlemoyer. QLoRA: Efficient finetuning of quantized LLMs. Advances in Neural Information Processing Systems , 2023.

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 on Learning Representations , 2021.

Timothy Dozat. Incorporating Nesterov momentum into Adam. In Proceedings of the International Conference on Learning Representations Workshop , 2016.

John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research , 12(61):2121-2159, 2011.

Sai Surya Duvvuri, Fnu Devvrit, Rohan Anil, Cho-Jui Hsieh, and Inderjit S Dhillon. Combining axes preconditioners through Kronecker approximation for deep learning. In Proceedings of the International Conference on Learning Representations , 2024.

Vladimir Feinberg, Xinyi Chen, Y. Jennifer Sun, Rohan Anil, and Elad Hazan. Sketchy: Memory-efficient adaptive regularization with frequent directions. Advances in Neural Information Processing Systems , 2023.

Elias Frantar, Eldar Kurtic, and Dan Alistarh. M-FAC: Efficient matrix-free approximations of second-order information. Advances in Neural Information Processing Systems , 2021.

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: Convolutionaugmented transformer for speech recognition. In Proceedings of the Conference of the International Speech Communication Association , 2020.

Chun-Hua Guo and Nicholas J. Higham. A Schur-Newton method for the matrix pth root and its inverse. SIAM Journal on Matrix Analysis and Applications , 28(3):788-804, 2006.

Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimization. In Proceedings of the International Conference on Machine Learning , 2018.

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.

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition , June 2016.

Roger A. Horn and Charles R. Johnson. Matrix Analysis . Cambridge university press, 2012.

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Proceedings of the International Conference on Learning Representations , 2015.

Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.

Ya Le and Xuan Yang. Tiny imagenet visual recognition challenge. CS 231N , 7(7):3, 2015.

Seung Hoon Lee, Seunghyun Lee, and Byung Cheol Song. Vision transformer for small-size datasets. arXiv preprint arXiv:2112.13492 , 2021.

Bingrui Li, Jianfei Chen, and Jun Zhu. Memory efficient optimizers with 4-bit states. Advances in Neural Information Processing Systems , 2023.

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.

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 IEEE/CVF International Conference on Computer Vision , October 2021.

Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In Proceedings of the International Conference on Learning Representations , 2019.

James Martens and Roger Grosse. Optimizing neural networks with Kronecker-factored approximate curvature. In Proceedings of the International Conference on Machine Learning , 2015.

Ning Qian. On the momentum term in gradient descent learning algorithms. Neural Networks , 12(1):145-151, 1999.

Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, and others. Language models are unsupervised multitask learners. OpenAI blog , 2019.

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.

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 FeiFei. ImageNet large scale visual recognition challenge. International Journal of Computer Vision , 115(3):211-252, 2015.

Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In Proceedings of the International Conference on Machine Learning , 2018.

Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. In Proceedings of the International Conference on Learning Representations , 2015.

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.09288 , 2023.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Ł ukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Information Processing Systems , 30, 2017.

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 on Artificial Intelligence , 2021.

Jui-Nan Yen, Sai Surya Duvvuri, Inderjit S. Dhillon, and Cho-Jui Hsieh. Block low-rank preconditioner with shared basis for stochastic optimization. Advances in Neural Information Processing Systems , 2023.

Hongwei Yong, Ying Sun, and Lei Zhang. A general regret bound of preconditioned gradient method for DNN training. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , June 2023.

Lin Zhang, Shaohuai Shi, and Bo Li. Eva: Practical second-order optimization with Kroneckervectorized approximation. In Proceedings of the International Conference on Learning Representations , 2023.

Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. GaLore: Memory-efficient LLM training by gradient low-rank projection. In Proceedings of the International Conference on Machine Learning , 2024.

Pan Zhou, Xingyu Xie, and Shuicheng Yan. Win: Weight-decay-integrated Nesterov acceleration for adaptive gradient algorithms. In Proceedings of the International Conference on Learning Representations , 2023.

A Implementation Details of Shampoo, CASPR, K-FAC and AdaBK

The implementation of 32-bit Shampoo used in our experiments is described in Algorithm 4. Our Pytorch implementation of Shampoo is partially based on the code provided by [2]. We implement CASPR by replacing ̂ G t = ̂ L t G t ̂ R t with J t = ̂ L t G t + G t ̂ R t ; ̂ G t = ̂ L t J t + J t ̂ R t in line 12 of Algorithm 4 and line 14 of Algorithm 3. We summarize the implementation of 32-bit K-FAC/AdaBK in Algorithm 5, where X t is the input feature and Y t is the output feature gradient. Both power iteration [5] and Schur-Newton iteration [17] are run for 10 iterations. Our implementation of 4-bit K-FAC/AdaBK is similar to 4-bit Shampoo (i.e., compressing L t , R t , ̂ L t , and ̂ R t ).

Algorithm 4 Practical 32-bit Shampoo

m × n

Input: initial parameter W 0 ∈ R , left preconditioner L 0 = ϵ I m , right preconditioner R 0 = ϵ I n , inverse root of left preconditioner ̂ L 0 = I m , inverse root of right preconditioner ̂ R 0 = I n , total number of steps T , interval of updating preconditioners T 1 , interval of updating inverse roots of preconditioners T 2 , exponential decay rate for preconditioners β ∈ (0 , 1) , first-order optimizer F , first-order optimizer state s 0 = 0 .

Output: final parameter W

T .

1: for t = 1 , 2 , . . . , T do

2: Receive loss function L t : R m × n ↦→ R and compute gradient G t = ∇L t ( W t )

3: if t % T 1 ≡ 0 then

4: L t = β L t -1 +(1 -β ) G t G T t ; R t = β R t -1 +(1 -β ) G T t G t

5: else

6: L t = L t -1 ; R t = R t -1

7: if t % T 2 ≡ 0 then

8: Compute maximum eigenvalues λ L max and λ R max of L t and R t by power iteration

9: Compute ̂ L t = ( L t + λ L max ϵ I m ) -1 / 4 and ̂ R t = ( R t + λ R max ϵ I n ) -1 / 4 by Schur-Newton iteration

10: else

11: ̂ L t = ̂ L t -1 ; ̂ R t = ̂ R t -1

12: ̂ G t = ̂ L t G t ̂ R t ; ˜ G t = ̂ G t ( ∥ G t ∥ F / ∥ ̂ G t ∥ F )

13: W t , s t = F ( W t -1 , s t -1 , ˜ G t )

Algorithm 5 Practical 32-bit K-FAC/AdaBK

Input: initial parameter W 0 ∈ R m × n , left preconditioner L 0 = 0 , right preconditioner R 0 = 0 , inverse root of left preconditioner ̂ L 0 = I m , inverse root of right preconditioner ̂ R 0 = I n , total number of steps T , interval of updating preconditioners T 1 , interval of updating inverse roots of preconditioners T 2 , ϵ , exponential decay rate for preconditioners β ∈ (0 , 1) , α = 1 for K-FAC / α = 2 for AdaBK, first-order optimizer F , first-order optimizer state s 0 = 0 .

Output: final parameter W T .

1: for t = 1 , 2 , . . . , T do

2: Receive loss function L t : R m × n ↦→ R and compute gradient G t = ∇L t ( W t )

3: Receive X t by forward propagation and Y t by backward propagation

4: if t % T 1 ≡ 0 then

5: L t = β L t -1 +(1 -β ) Y t Y T t ; R t = β R t -1 +(1 -β ) X t X T t

6: else

7: L t = L t -1 ; R t = R t -1

8: if t % T 2 ≡ 0 then

9: Compute maximum eigenvalues λ L max and λ R max of L t and R t by power iteration

10: Compute ̂ L t = ( L t + λ L max ϵ I m ) -1 /α and ̂ R t = ( R t + λ R max ϵ I n ) -1 /α by Schur-Newton iteration

11: else

12: ̂ L t = ̂ L t -1 ; ̂ R t = ̂ R t -1

13: ̂ G t = ̂ L t G t ̂ R t ; ˜ G t = ̂ G t ( ∥ G t ∥ F / ∥ ̂ G t ∥ F )

14: W t , s t = F ( W t -1 , s t -1 , ˜ G t )

0

1

2

3

4

5

6

7

Index

1.00

0.75

0.50

0.25

0.00

0.25

0.50

0.75

1.00

Value

DT

Linear-2

(a) 3-bit quantization

0

2

4

6

8

10

12

14

Index

1.00

0.75

0.50

0.25

0.00

0.25

0.50

0.75

1.00

Value

DT

Linear-2

(b) 4-bit quantization

B Randomized SVD Method

Given an initial matrix P 0 ∈ R n × n , randomized SVD method computes the eigenvector matrix of a PD matrix A ∈ R n × n by iterating

where QR ( X ) denotes the QR decomposition of matrix X , returning an orthogonal matrix. Since we can initialize P 0 with the previous result (e.g., V in Algorithm 1), only a few iterations are enough to obtain an accurate estimation in practice. In our experiments, we iterate (4) once for Shampoo/CASPR, and iterate (4) twice for K-FAC/AdaBK.

C Quantization Mappings

We present the constructions of different quantization mappings in b -bit quantizers ( R in Q ). See Figure 5 for the illustration of them. Note that T b = { 0 , 1 , . . . , 2 b -1 } .

Dynamic tree (DT) quantization for b -bit quantization maps T b onto { 0 , 1 } ∪ G , where G is a set of numbers with the following properties: the number in G looks like ± q k × 10 -E , where a) b = 2 + E + F , where E,F are integers; b) q k = ( p k + p k +1 ) / 2 , where k ∈ { 0 , . . . , 2 F -1 } ; c) p j = 0 . 9 j/ 2 F +0 . 1 , where j ∈ { 0 , . . . , 2 F } . For 4-bit quantization, DT quantization maps T 4 onto {-0.8875, -0.6625, -0.4375, -0.2125, -0.0775, -0.0325, -0.0055, 0.0000, 0.0055, 0.0325, 0.0775, 0.2125, 0.4375, 0.6625, 0.8875, 1.0000}. For 3-bit quantization, DT quantization maps T 3 onto {-0.7750, -0.3250, -0.0550, 0.0000, 0.0550, 0.3250, 0.7750, 1.0000}.

For 4-bit quantization, linear square (Linear-2) quantization maps T 4 onto {-1.0000, -0.7511, -0.5378, -0.3600, -0.2178, -0.1111, -0.0400, 0.0000, 0.0044, 0.0400, 0.1111, 0.2178, 0.3600, 0.5378, 0.7511, 1.0000}. For 3-bit quantization, Linear-2 quantization maps T 3 onto {-1.0000, -0.5102, -0.1837, 0.0000, 0.0204, 0.1837, 0.5102, 1.0000}.

D Quantization Error Analyses

We present more quantization error analyses of the preconditioners. Recall that we define two kinds of quantization errors in mapping f of transformation g at A ∈ R m × n (in short errors in f ( A ) of g ) in Subsection 3.1. Here we extend them as follows: define the normwise relative error (NRE) in f of ( g 1 , g 2 ) at A as

and the angle error (AE) in f of ( g 1 , g 2 ) at A as

D.1 Static Analysis

Table 5 is an extension of Table 1 for Bit=4. Since the diagonal elements of A -1 / 4 are usually much larger than its non-diagonal elements where A is a PD matrix, we further consider the quantization errors in f ( A )= A -1 / 4 -Diag(diag( A -1 / 4 )) at 4-bit precision as shown in Table 6. Table 7 shows the quantization errors at 8-bit precision.

A large condition number of a PD matrix A is indispensable for the superiority of quantizing U over quantizing A , where U is the eigenvector matrix of A . We consider contracting the singular value distribution of A = A 1 with SVD U Diag( λ ) U T used in Table 5 by mapping each singular value λ of A to h ( λ ) = τ ( λ -λ A min ) + λ A min , where λ A min is the minimum singular value of A and τ > 0 is the contraction coefficient. Figure 6 shows 4-bit quantization errors in A -1 / 4 or A -1 / 4 -Diag(diag( A -1 / 4 )) of quantizing U or A at A = U Diag( h ( λ )) U T .

20

16

12

8

4

0

Contraction Coefficient

0.0

0.1

0.2

0.3

0.4

0.5

0.6

Normwise Relative Error

quan

A

quan

U

20

16

12

8

4

0

Contraction Coefficient

0.0

2.5

5.0

7.5

10.0

12.5

15.0

17.5

)

quan

A

quan

U

(a) Errors in

f

(

A

)=

A

-

1

/

4

20

16

12

8

4

0

Contraction Coefficient

0.0

0.2

0.4

0.6

0.8

1.0

Normwise Relative Error

quan

A

quan

U

20

16

12

8

4

0

Contraction Coefficient

10

20

30

40

50

60

70

80

)

quan

A

quan

U

(b) Errors in

f

(

A

)=

A

-

1

/

4

-

Diag(diag(

A

-

1

/

4

))

D.2 Dynamic Analysis

We define the normwise relative error (NRE) and angle error (AE) of B deviating from A as

Consider Shampoo using 4-bit preconditioners for parameter updates, but also recording 32-bit preconditioners at the same time. We extract the left preconditioners L 4 and L 32 ∈ R 1200 × 1200 of a specific model parameter block W ∈ R 1200 × 768 every 8000 steps in the Swin-Tiny training on CIFAR-100 with AdamW+Shampoo. Here L 4 is a decompressed 4-bit preconditioner, and L 32 is a 32-bit preconditioner.

Figure 7 shows the quantization errors during training. For naive 4-bit Shampoo, L -1 / 4 32 and L -1 / 4 4 are computed by Schur-Newton iteration used in Algorithm 4 where ϵ =10 -4 . For our 4-bit Shampoo, L -1 / 4 32 is computed by Schur-Newton iteration used in Algorithm 4 where ϵ =10 -4 , and L -1 / 4 4 is computed by Algorithm 2 without quantization where ϵ =10 -4 , t 2 =4 . We find that ϵ =10 -6 for Algorithm 2 used in our main experiments though is effective, yet it can cause a large numerical instability in the later stage of training (see Figure 8).

2

4

6

8

10

12

14

Step (×8000)

0.2

0.3

0.4

0.5

0.6

0.7

0.8

Normwise Relative Error

4-bit Shampoo (naive)

4-bit Shampoo (our)

2

4

6

8

10

12

14

Step (×8000)

10

12

14

16

18

20

22

)

4-bit Shampoo (naive)

4-bit Shampoo (our)

(a) Errors in

L

-

1

/

4

4

deviating from

L

-

1

/

4

32

2

4

6

8

10

12

14

Step (×8000)

0.70

0.75

0.80

0.85

0.90

0.95

1.00

Normwise Relative Error

4-bit Shampoo (naive)

4-bit Shampoo (our)

2

4

6

8

10

12

14

Step (×8000)

40

45

50

55

60

65

)

4-bit Shampoo (naive)

4-bit Shampoo (our)

(b)

Errors in

L

-

1

/

4

4

-

Diag(diag(

L

-

1

/

4

4

))

deviating

from

L

-

1

/

4

32

-

Diag(diag(

L

-

1

/

4

32

))

2

4

6

8

10

12

14

Step (×8000)

0.2

0.3

0.4

0.5

0.6

0.7

0.8

Normwise Relative Error

4-bit Shampoo (naive)

4-bit Shampoo (our)

2

4

6

8

10

12

14

Step (×8000)

12

14

16

18

20

22

)

4-bit Shampoo (naive)

4-bit Shampoo (our)

(a) Errors in

L

-

1

/

4

4

deviating from

L

-

1

/

4

32

2

4

6

8

10

12

14

Step (×8000)

0.8

1.0

1.2

1.4

1.6

Normwise Relative Error

4-bit Shampoo (naive)

4-bit Shampoo (our)

2

4

6

8

10

12

14

Step (×8000)

45

50

55

60

65

)

4-bit Shampoo (naive)

4-bit Shampoo (our)

(b)

Errors in

L

-

1

/

4

4

-

Diag(diag(

L

-

1

/

4

4

))

deviating

from

L

-

1

/

4

32

-

Diag(diag(

L

-

1

/

4

32

))

E Convergence Analysis

More notations. Given a symmetric real matrix A , A ⪰ 0 means that A is positive semidefinite (PSD), and A ≻ 0 means that A is positive definite (PD). Assume that symmetric matrices A and B are symmetric, the notations A ⪰ B and A ≻ B mean that A -B ⪰ 0 and A -B ≻ 0 respectively. Let A be a PSD matrix and s ∈ R , we define A s = U Λ s U T , where U Λ U T is the Singular Value Decomposition (SVD) of A . The Mahalanobis norm of a vector x induced by a PD matrix A is ∥ x ∥ A = √ x T Ax . The dual norm of ∥ · ∥ A is denoted by ∥ · ∥ ∗ A , where ∥ x ∥ ∗ A = √ x T A -1 x . The spectral norm of matrix A is ∥ A ∥ 2 = sup x = 0 {∥ Ax ∥ 2 / ∥ x ∥ 2 } . A ⊗ B means the (right) Kronecker product of matrices A and B . vec( A ) means the vectorization (stacking the rows) of A .

̸

Algorithm 6 Perturbed Shampoo in the matrix case

Input: W 0 ∈ R m × n , L 0 = 0 m × m , R 0 = 0 n × n , ρ 0 = 0 , µ 0 = 0 .

1: for t = 1 , . . . , T do

2: Receive loss function: f t : R m × n → R

3: Compute gradient: G t = ∇ f t ( W t )

4: Update preconditioners: J t = L t -1 + G t G T t ; K t = R t -1 + G T t G t

5: Perturb preconditioners: L t = g ( J t ); R t = g ( K t )

6: Accumulate errors: ρ t = ρ t -1 + ∥ J t -L t ∥ 2 ; µ t = µ t -1 + ∥ K t -R t ∥ 2

7: Update parameters: W t +1 = W t -η (( ϵ + ρ t ) I m + L t ) -1 / 4 G t (( ϵ + µ t ) I n + R t ) -1 / 4

We consider quantization as a perturbation and present the perturbed Shampoo in Algorithm 6 for convergence analysis. The regret bound of the perturbed Shampoo can be found in Theorem 1. Complete proofs can be found in Appendix F. We first introduce some basic technical tools, and the details of them are in [18, 21].

Lemma 3. Let A , A ′ , B , B ′ be matrices of appropriate dimensions, and u , v be two column vectors. The following properties hold:

( A ⊗ B ) T = ( A T ⊗ B T );

If A , B ⪰ 0 and s ∈ R , then ( A ⊗ B ) s = ( A s ⊗ B s );

If A ⪰ A ′ and B ⪰ B ′ , then A ⊗ B ⪰ A ′ ⊗ B ′ ;

tr( AB ) = tr( A )tr( B );

vec( uv T ) = u ⊗ v .

Lemma 4. Let G ∈ R m × n , L ∈ R m × m , R ∈ R n × n , then it holds that

Lemma 5. Assume that 0 ⪯ X i ⪯ Y i for i = 1 , . . . , n . Assume further that all X i commute with each other and all Y i commute with each other. Let α 1 , . . . , α n ≥ 0 such that ∑ n i =1 α i = 1 , then

Lemma 6. Let 0 ≤ α ≤ 1 and 0 ⪯ X ⪯ Y , then X α ⪯ Y α .

Lemma 7. Let A ≻ 0 and B ≻ 0 , then it holds that A ⪰ B if and only if B -1 ⪰ A -1 .

Lemma 8 (von Neumann) . Let A , B ∈ R m × n and q = min { m,n } . Let σ 1 ( A ) ≥ · · · ≥ σ q ( A ) and σ 1 ( B ) ≥ · · · ≥ σ q ( B ) denote the non-increasingly ordered singular values of A and B , respectively. Then

Lemma 9. Assume that function f t is continuously differentiable and convex on R d , and matrix H t ≻ 0 for t = 1 , . . . , T . Given w 0 ∈ R d , η > 0 , define w t +1 = w t -η H -1 t g t , where g t = ∇ f t ( w t ) . Then for any w ∗ ∈ R d , we have

Lemma 10. Let g 1 , . . . , g T be a sequence of vectors. For ρ > 0 , define ̂ H t = ( ρ I + ∑ t s =1 g s g T s ) 1 / 2 . Then we have

Lemma11. Assume that G 1 , . . . , G T ∈ R m × n are matrices of rank at most r . Let s for t = 1 , . . . , T . Then for any ϵ ≥ 0 ,

The key to the convergence proof of Algorithm 6 is forming a PD matrix sequence { H i } T i =1 , which satisfies 0 ≺ H 1 ⪯ · · · ⪯ H T . To achieve it, we gives the following lemma extended from Lemma 2 in the Appendix of [40].

Lemma 12. Let { X t } t = T t =1 be a sequence of symmetric matrices, and A t = ∑ t s =1 X s , where t = 1 , . . . , T . Suppose we have two sequences of symmetric matrices { Y t } t = T t =1 , { Z t } t = T t =0 , and a sequence real numbers { ρ t } t = T t =0 satisfying

Define B t = ρ t I + Z t , where I denotes the identity matrix. Then for t = 1 , . . . , T , we have

Theorem 1. Assume that the gradients G 1 , . . . , G T ∈ R m × n are matrices of rank at most r . Then for any W ∗ ∈ R m × n and ϵ > 0 , if η = D/ √ 2 r , the regret of Algorithm 6 is bounded as follows,

Though we get a convergence guarantee of Algorithm 6, the upper bound given by Theorem 1 is very slack, since 2 1 / 4 mρ 1 / 4 T is about the same as tr( ˜ L 1 / 4 T ) for 4-bit quantization schemes in practice.

F Proofs

Lemma 1. Let A be a PD matrix whose SVD is U Λ U T , where U =[ u i ] is an orthogonal matrix and Λ =diag([ λ i ] T ) is a diagonal matrix. Given a perturbation ∆ U =[∆ u i ] and s ∈ R , we define B :=( U Λ U T ) s and ∆ B :=(( U +∆ U ) Λ ( U +∆ U ) T ) s -B .

(1) If U +∆ U is orthogonal and there exists α ∈ R such that ∥ ∆ u i ∥ 2 ≤ α , then

(2) If U +∆ U is orthogonal and there exists β ∈ R such that ⟨ u i , u i +∆ u i ⟩ ≥ 1 -β ≥ 0 , then

Proof. (1) Since U and U +∆ U are orthogonal, we have

by definition. This leads to

The Frobenius norm satisfies the triangle inequality and is orthogonality invariant. Hence,

(2) Similar to (1), we have

From ⟨ u i , u i +∆ u i ⟩ ≥ 1 -β ≥ 0 , we get 0 ≥ ⟨ u i , ∆ u i ⟩ ≥ -β ≥ -1 because

holds due to the orthogonality of U and U +∆ U . Hence,

Therefore, we have

The proof is completed.

Lemma 2. Let A be a PD matrix of order m + n whose SVD is U Λ U T , where m,n ∈ N + , n = lm , U = [ u i ] is an orthogonal matrix and Λ = diag([ λ i ] T ) is a diagonal matrix. Assume that Λ = diag([ cλ 1 T m × 1 , λ 1 T n × 1 ] T ) , c ≥ 1 , and λ > 0 . Given a perturbation ∆ Λ = diag([ 0 T m × 1 , ∆ λ T n × 1 ] T ) and s ∈ R , we define B :=( U Λ U T ) s and ∆ B :=( U ( Λ +∆ Λ ) U T ) s -B .

If ∆ λ n × 1 = ( k -1) λ 1 n × 1 where k > 0 , then

Moreover, h 1 ( s, l ) decreases monotonically with s over ( -∞ , 0) and increases monotonically with l over (0 , + ∞ ) .

If ∆ λ n × 1 = ( tc -1) λ 1 n × 1 where t > 0 , then

Moreover, h 2 ( l ) decreases monotonically with l over (0 , ( c/t ) s ] and increases monotonically with l over (( c/t ) s , + ∞ ) .

If ∆ λ n × 1 = ( tc -1) λ 1 n × 1 where k = tc > 0 and l = ( c/t ) s , then

Proof. (1) Since U is orthogonal, we have

Hence,

It is easy to check that h 1 increases monotonically with l over (0 , + ∞ ) . To prove h 1 decreases monotonically with s over ( -∞ , 0) , define

Consider the derivative of g 1

If s < 0 and k > 1 , then k s -1 < 0 , k s ln k > 0 leading to g ′ 1 ( s ) < 0 since c ≥ 0 ; Similarly, if s < 0 and 0 < k ≤ 1 , then k s -1 ≥ 0 , k s ln k ≤ 0 leading to g ′ 1 ( s ) ≤ 0 . Thus g 1 ( s ) is a monotonically decreasing function for s < 0 , which implies that h 1 decreases monotonically with s over ( -∞ , 0) .

(2) Similar to (1), we have

Besides,

Hence, we get

To prove h 2 decreases monotonically with l over (0 , ( c/t ) s ] and increases monotonically with l over (( c/t ) s , + ∞ ) , we define

whose monotonicity is equivalent to that of h 2 for l > 0 . Consider the derivative of g 2

̸

If s = 0 or tc = 1 , then g 2 ( l ) ≡ 1 . If s = 0 and tc = 1 , then ( t s -t 2 s c s ) 2 > 0 , ( c s -t s c 2 s ) 2 > 0 . In this case, let g ′ 2 ( l ) = 0 , we get

̸

which implies that l = ( c/t ) s . It is easy to see that g 2 decreases monotonically with l over (0 , ( c/t ) s ] and increases monotonically with l over (( c/t ) s , + ∞ ) .

(3) According to (1)(2), we can easily get

The proof is completed.

Proposition 1. Let A be a PD matrix of order m + n whose SVD is U Λ U T , where m,n ∈ N + , n = lm , U =[ u i ] is an orthogonal matrix, Λ =diag([ cλ 1 T m × 1 , λ 1 T n × 1 ] T ) , c ≥ 1000 , and λ > 0 . Given ∆ U =[∆ u i ] , ∆ Λ =diag([ 0 T m × 1 , ∆ λ T n × 1 ] T ) , and s ≤-0 . 25 , we define B := ( U Λ U T ) s , B 1 := (( U + ∆ U ) Λ ( U +∆ U ) T ) s , and B 2 := ( U ( Λ +∆ Λ ) U T ) s . If U + ∆ U is orthogonal, ∥ ∆ u i ∥ 2 ≤ 0 . 1 , ⟨ u i , ∆ u i ⟩≥-0 . 005 , ∆ λ n × 1 =(0 . 02 c -1) λ 1 n × 1 , and l =( c/ 0 . 02) s , then

Proof. According to Lemma 1, we have

On the other hand, from Lemma 2(3), we get

where x = (0 . 02 c ) s ∈ (0 , 20 -1 / 4 ] . It is easy to verify that f 1 decreases monotonically and f 2 increases monotonically for 0 < x < 1 . Hence

The proof is completed.

Lemma 12. Let { X t } t = T t =1 be a sequence of symmetric matrices, and A t = ∑ t s =1 X s , where t = 1 , . . . , T . Suppose we have two sequences of symmetric matrices { Y t } t = T t =1 , { Z t } t = T t =0 , and a sequence real numbers { ρ t } t = T t =0 satisfying

Define B t = ρ t I + Z t , where I denotes the identity matrix. Then for t = 1 , . . . , T , we have

Proof. Note that for any symmetric matrix S , it holds that ∥ S ∥ 2 I ⪰ S . Then we have

Adding ρ t -1 I on both sides, we get

Hence

On the other hand, we have

Adding ρ t I on both sides, we get

Hence

The proof is completed.

Theorem 1. Assume that the gradients G 1 , . . . , G T ∈ R m × n are matrices of rank at most r . Then for any W ∗ ∈ R m × n and ϵ > 0 , if η = D/ √ 2 r , the regret of Algorithm 6 is bounded as follows,

where D = max t ∈ [ T ] ∥ W t -W ∗ ∥ F , ˜ L t = ϵ I m + ∑ T t =1 G t G T t , and ˜ R t = ϵ I n + ∑ T t =1 G T t G t

Proof. Define ˆ L t = ( ϵ + ρ t ) I m + L t , ˆ R t = ( ϵ + µ t ) I n + R t . According to Lemma 12, ˆ L t and ˆ R t are positive definite. Recall the update performed in Algorithm 6,

For t > 0 , let H t = ˆ L 1 / 4 t ⊗ ˆ R 1 / 4 t , g t = vec( G t ) and w t = vec( W t ) . Due to Lemma 3(3) and Lemma 4, we have

Lemma 12 implies 0 ≺ ˆ L 1 ⪯ · · · ⪯ ˆ L T , 0 ≺ ˆ R 1 ⪯ · · · ⪯ ˆ R T . Thus, according to Lemma 3(3)(4) and Lemma 6, we get

Let H 0 = 0 . By invoking Lemma 9 and Lemma 8, we obtain the regret bound

Define ̂ H t = ( rϵ I + ∑ t s =1 g s g T s ) 1 / 2 . Lemma 11 and Lemma 12 imply that

Using Lemma 7 and Lemma 10 along with the above equation, we obtain

Consequently, using Lemma 3(5) and Lemma 12, we get the desired regret bound

by choosing η = D/ √ 2 r . The proof is completed.

G Experimental Details

We use one RTX3060Ti GPU under the PyTorch 2.0.1+CUDA11.8 framework for DNN training on the CIFAR-100 and Tiny-ImageNet datasets, use one A800 GPU under the PyTorch 2.0.1+CUDA11.7 framework for DNN training on the ImageNet-1k and C4 datasets, and use two NVIDIA L40S GPUs under the PyTorch 2.0.1+CUDA11.8 framework for DNN training on the OWT dataset. To obtain the total peak memory consumption per GPU, we call "torch.cuda.max_memory_allocated".

We set "torch.backends.cudnn.benchmark" to "False" for all the experiments, except when training ViT-Base/32 on the ImageNet-1k dataset. We report the total memory consumption instead of the memory consumption of the second-order optimizer. This total memory includes data, model parameters, activations, gradients, states forming the preconditioners and their inverse roots, states for the used first-order optimizer, and memory fragments. Our focus lies in quantizing the states for constructing preconditioners and their inverse roots, which are approximately 7x smaller for 4-bit Shampoo compared to 32-bit Shampoo. Because the block size is 64, its maximum value should be calculated every 64 elements and saved as a 32-bit value, resulting in an additional overhead of 0.5 bits ( 32 / 64 ). Consequently, the memory savings are approximately 7 times, calculated as 32 / (4+0 . 5) . In the future, we may adopt double quantization [9] to further reduce memory consumption.

For SGDM, Adagrad or AdamW used in second-order optimizers, we use 32-bit optimizer states on image classification tasks and 16-bit optimizer states on natural language modeling tasks by default. For SGDM, we set the momentum to 0.9 and use an initial learning rate of 0.1. For Adagrad, we set ϵ = 10 -10 and use an initial learning rate of 0.01. For AdamW, we set β 1 = 0 . 9 , β 2 = 0 . 999 , and ϵ = 10 -8 and use an initial learning rate of 0.001. For quantization settings, we employ block-wise normalization with a block size of 64 and linear square quantization by default. Matrices with a size smaller than 4096 will not be quantized. For Shampoo and CASPR, we use ϵ = 10 -6 , β = 0 . 95 and t 1 = 1 , t 2 = 4 by default. Shampoo and CASPR precondition blocks from large matrices and the maximum order of a preconditioner is 10000 for 130M LLAMA-2 and is 1200 for other models. For training loss, we use cross-entropy loss. For image classification tasks, automatic mixed precision is enabled except for training transformers on the CIFAR-100 and Tiny-ImageNet datasets.

Settings on training CNNs on CIFAR-100 or Tiny-ImageNet. Minibatch size is set to 128. Weight decay is 0.0005. Data augmentation includes random crop and horizontal flip. For Shampoo, we set T 1 = 100 and T 2 = 500 . In Section 5, we run SGDM for 300 epochs and SGDM+Shampoo for 200 epochs on the CIFAR-100 dataset. We run SGDM for 150 epochs and SGDM+Shampoo for 100 epochs on the Tiny-ImageNet dataset. We adopt the multi-step learning rate schedule (the learning rate is multiplied by 0.1 for every 30% epochs with a linear warmup at the first 5 epochs).

Settings on training transformers on CIFAR-100 or Tiny-ImageNet. We set a patch size of 4 for ViT-small on the CIFAR-100 dataset, and a patch size of 8 for ViT-small on the Tiny-ImageNet dataset. For training Swin-Tiny on the CIFAR-100 dataset, we use a patch size of 2 and window size of 4. For training Swin-Tiny on the Tiny-ImageNet dataset, we use a patch size of 4 and window

size of 7. Minibatch size is set to 128. We run Adagrad/AdamW/NadamW for 150 epochs and Adagrad/AdamW+Shampoo for 100 epochs. Weight decay is 0.0005 for Adagrad, and is 0.05 for AdamW/NadamW. We use the cosine learning rate schedule. Data augmentation follows the source code in [25]. For Shampoo, we set T 1 = 100 and T 2 = 500 . With the exception of certain optimizer settings, the configurations used for ablation studies are identical to those outlined above.

Settings on training ResNet50 on ImageNet-1k. 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 is 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 T 1 = 200 and T 2 = 1000 .

Settings on training ViT-Base/32 on ImageNet-1k. We run AdamW for 150 epochs and AdamW+Shampoo for 120 epochs. Minibatch size is set to 512. Weight decay is 0.05. We use the cosine learning rate schedule. Data augmentation follows the configuration for training ViT-Base/16 in [44], excluding repeated augmentation. For Shampoo, we set T 1 = 200 and T 2 = 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 T 1 = 200 and T 2 = 200 . For our 4-bit Shampoo, we use Schur-Newton iteration used in Algorithm 4 to compute the inverse root of a preconditioner for training stability.

Settings on training LLAMA-2 on C4. We run AdamW with 10% warmup steps. Total batch size is 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 bfloat16. Weight decay is 0. For Shampoo, we set T 1 = 200 and T 2 = 200 .

Settings on K-FAC and AdaBK. K-FAC/AdaBK preconditions layers without limiting the size of a preconditioner. We set β = 0 . 9 , T 1 = 200 , and T 2 = 2000 . We use ϵ = 0 . 1 for K-FAC and ϵ = 0 . 001 for AdaBK. For 4-bit K-FAC/AdaBK, we set t 1 = 0 and t 2 = 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, we set lr=1.0, weight_decay=0.0005 and warmup_steps=2000. For AdamWScheduleFree, we set lr=0.0025, weight_decay=0.05 and warmup_steps=10000.

Settings on M-FAC. We use the code from [15] and set ngrads=32, damp=0.1. The other hyperparameter 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 8 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.

We also provide the results of training ResNet34 and Swin-Tiny on CIFAR-100 with schedulefree approach [6] in Table 9. From it one can see that AdamWScheduleFree achieves comparable performance to AdamW with cosine decay, while SGDScheduleFree underperforms compared to SGDM. We observe that this schedule-free algorithm shows rapid improvements in training and test accuracy during the early training stages, but may fail to achieve a higher test accuracy ultimately (see Figure 9). Anyway, these methods are still worse than our AdamW+4-bit Shampoo.

0

50

100

150

200

250

300

Epoch

50

55

60

65

70

75

80

Test Accuracy (%)

ResNet34 on CIFAR-100

SGDM

SGDScheduleFree

0

20

40

60

80

100

120

140

Epoch

50

55

60

65

70

75

Test Accuracy (%)

Swin-Tiny on CIFAR-100

AdamW

AdamWScheduleFree

More optimizers. Table 10 shows results of training Swin-Tiny on CIFAR-100 with NadamW, Adagrad and Adagrad+Shampoo. One can see that Adagrad+4-bit Shampoo converges faster than Adagrad with ignorable extra memory overhead, and also has higher test accuracy. Besides, though NadamW [11] is slightly better than AdamW, it is still worse than our AdamW+4-bit Shampoo.

M-FAC [15] is a matrix-free method computing inverse-Hessian vector products with many gradient copies. It is not memory-efficient for M-FAC to maintain m dense gradient copies ( m = 1024 in its official code). Table 11 shows that both SGDM+32-bit Shampoo and SGDM+4-bit Shampoo enjoy much higher efficiency than M-FAC ( m = 32 ) for training ResNet34 on CIFAR-100, and enjoy higher test accuracy. EVA [42] is a rank-one second-order optimizer and is memory-efficient. We train ResNet34 on CIFAR-100 with SGDM+EVA, but despite extensive hyper-parameter tuning, we fail to achieve acceleration over SGDM. Instead, we cite EVA's result of training VGG-19 on CIFAR-100 for 200 epochs (see Table 2 in [42]). The test accuracies of SGDM+EVA and SGDM+Shampoo are 73% and 74.5%, respectively.

260

280

300

320

340

360

380

Wall-clock Time (min)

3.18

3.19

3.20

3.21

3.22

3.23

3.24

Validation Loss

LLAMA-130M on C4

AdamW

AdamW+32-bit Shampoo

AdamW+4-bit Shampoo (naive)

AdamW+4-bit Shampoo (our)

1900

2000

2100

2200

2300

2400

Wall-clock Time (min)

2.93

2.94

2.95

2.96

2.97

2.98

Validation Loss

GPT2-124M on OWT

AdamW

AdamW+32-bit Shampoo

AdamW+4-bit Shampoo (naive)

AdamW+4-bit Shampoo (our)

H.2 Natural Language Modeling

Models, datasets, and hyperparameters. We train 124M GPT-2 [32] for 60k steps on the OpenWebText (OWT) dataset * following the nanoGPT codebase † with two NVIDIA L40S GPUs, and train 130M LLAMA-2 [37] for 20k steps and 350M LLAMA-2 for 60k steps on the C4 dataset [33] following [43] with one A800 GPU. See Appendix G for experimental details.

Main results. We show the performance, wall-clock time, and memory cost in Table 12, and the validation loss curves in Figure 10. As with the vision tasks, our AdamW+4-bit Shampoo consistently outperformed AdamW and naive AdamW+4-bit Shampoo in terms of performance, and AdamW+32-bit Shampoo in terms of memory usage.

Memory efficiency. Wefurther check the memory usage by increasing token batch size for a language model, which is calculated as the batch size multiplied by the context length (see [43]). To train LLAMA2-7B on the C4 dataset using a single A800 GPU (with a maximum memory of 81,920

MB), we set the context length to 256 and then determine the maximum batch size allowed by each optimizer. For Shampoo, the maximum order of a preconditioner for training LLAMA2-7B is 2048. In all experiments, gradient checkpointing is enabled. Table 13 summarizes the evaluation results. By comparison, the 32-bit Shampoo runs out of memory with a batch size of 2, while our 4-bit Shampoo supports a batch size of 64 for standard training and only encounters memory issues at a batch size of 128. These results clearly demonstrate that our 4-bit Shampoo significantly conserves memory compared to the 32-bit version.

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 contributions 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 to include aspirational goals as motivation as long as it is clear that these goals are 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 no 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 violations 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 be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was 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 fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in 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 to 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-referenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We present the implementation details of all the experiments in the Appendix. Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to 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, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

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, as described in supplemental material?

Answer: [Yes]

Justification: We use publicly available datasets and will release our source code. Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/public/ guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

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:// nips.cc/public/guides/CodeSubmissionPolicy ) 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 scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are 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.

6. 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 to 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 a 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 are not reported because it would be too computationally expensive.

Guidelines:

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, at 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 a 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 or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 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 to reproduce the experiments?

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.

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 experimental 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 or 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 https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We review the NeurIPS Code of Ethics and our paper conforms it. Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: We discuss both potential positive societal impacts and negative societal impacts at the end of the paper.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: We present 4-bit Shampoo for memory efficient training of deep models. It poses no such risks.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We properly mention all the existing assets.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

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 dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

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 crowdsourcing 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 to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to 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 nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If 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 NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nQuantization is applied to eigen matrix instead of directly to precond. The eigen matrix is then orthonormalized via  Björck orthonormalization to orthogonalize V.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nCould 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.\n\nQUESTIONS:\nWould 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.\n\nLIMITATIONS:\nAs 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-06-13"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThe method should help in optimization of large networks, but no results are presented on large networks. \nThe presentation has some minor issues, noted below.\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-08"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. **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\n\nQUESTIONS:\nPlease refer to the section of weaknesses.\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-12"}

comments: 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.

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.
- The experiments span various datasets and models, consistently achieving performance that matches that of the full-precision Shampoo.
- 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.

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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2024-07-14