4-bit Shampoo for Memory-Efficient Network Training

Thank you for the insightful and positive comments. In the following, we provide our point-by-point response and hope our response helps address your concerns. We also look forward to the subsequent discussion which further helps to solve the current issues.

1) Add experiments on LLMs & 124M GPT2. Thanks. Due to limited computing resources and the short rebuttal period of 7 days, we trained medium-sized language models, including 124M GPT-2 for 20k steps on the OpenWebText (OWT) dataset using code from nanoGPT (https://github.com/karpathy/nanoGPT), and LLAMA-2 130M for 20k and 350M for 60k steps on the C4 dataset following [1]. For these experiments, we adhered to the exact settings provided in the corresponding papers or GitHub repositories to ensure a fair comparison. The results are reported in Table 2 and Figure 1 of the rebuttal PDF. As with the vision tasks in the manuscript, our AdamW+4-bit Shampoo consistently outperformed AdamW and AdamW+4-bit Shampoo (naive) in terms of performance, and AdamW+32-bit Shampoo in terms of memory usage. Since our algorithm design is not task-specific, we believe its improvements should be consistently achievable on other tasks as well which is also validated by the results on both the computer vision tasks in the manuscript and the natural language processing tasks here.

For LLMs, we have compared our quantization with the naive quantization approach in Table 2 and Figure 1 of the rebuttal PDF. Notably, our 4-bit Shampoo with our quantization outperforms the naive 4-bit Shampoo with naive quantization in terms of validation loss with almost the same training times and memory usage. This demonstrates that the performance gain of our quantization is task-independent.

2) More proper learning rate decay for ResNet. Thanks. Due to time limitation, we use cosine learning rate decay (initial lr=0.1) to train ResNet34 on CIFAR-100. Experimental results can be found in the table below. By comparison, SGDM+Shampoo still converges faster than SGDM, and have slightly better test performance.

3) Initial learning rate for ViT ImageNet 1-k. Thanks. We use the default initial learning rate provided in [2] to train ViT-Base. Although there may be a better learning rate, the performance reported in our paper is relatively high and reasonable. For ViT-Base/32 on ImageNet 1-k, Table 5 in [3] reports 73.38% accuracy for training 300 epochs, while ours for 150 epochs is 72.87% accuracy.

4) Comparison with NadamW and schedule free optimization. Per your suggestion, we provide the results of training Swin-Tiny on CIFAR-100 in the table below. We run AdamW, NadamW, and AdamWScheduleFree for 150 epochs and AdamW+Shampoo for 100 epochs. For NadamW, we use its implementation in Timm library, using the same hyperparameters as AdamW. For schedule free optimization, we use the code from [4] to train Swin-Tiny, and set lr=0.0025 (default), weight_decay=0.05, warmup_steps=10000.

One can see that though improving AdamW, these methods are still worse than our AdamW+4-bit Shampoo.

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

[2] Pan Zhou, Xingyu Xie, and Shuicheng Yan. Win: Weight-decay-integrated Nesterov acceleration for adaptive gradient algorithms. In ICLR, 2023.

[3] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In ICLR, 2021.

[4] The Road Less Scheduled. arXiv preprint arXiv:2405.15682, 2024.

Thank you for the insightful and positive comments. In the following, we provide our point-by-point response and hope our response helps address your concerns. We also look forward to the subsequent discussion which further helps to solve the current issues.

1) Training of large-scale models. Due to the limited computing resources and rebuttal period, we cannot afford sufficient pre-training of very large models. Instead, we train medium-sized language models, including 124M GPT-2 for 20k steps on the OpenWebText (OWT) dataset using the code from https://github.com/karpathy/nanoGPT, and LLAMA-2 130M for 20k and 350M for 60k steps on the C4 dataset following [1]. For these experiments, we adhered to the exact settings provided in the corresponding papers or GitHub repositories to ensure a fair comparison. The results are reported in Table 2 and Figure 1 of the rebuttal PDF. Similar to the vision tasks in the manuscript, our AdamW+4-bit Shampoo consistently outperformed AdamW and AdamW+4-bit Shampoo (naive) in terms of performance, and AdamW+32-bit Shampoo in terms of memory usage.

Additionally, in the manuscript, we have evaluated six models (including VGG, ResNet, ViT, and Swin) on three datasets: CIFAR100, Tiny-ImageNet, and the large-scale ImageNet dataset. These diverse experiments sufficiently demonstrate the superiority of our 4-bit Shampoo over the vanilla 32-bit version.

Finally, our algorithm design does not rely on any specific properties of the tasks, ensuring that the performance improvements are general and transferable. So we believe these improvements should be consistently achievable on other tasks as well. This is validated by the results on both the computer vision tasks in the manuscript and the natural language processing tasks presented in the rebuttal PDF.

2) Add a row for 8-bit A in Table 1. Per your suggestion, we provide the quantization errors of 8-bit quantization schemes in Table 1 of the rebuttal PDF. Comparing Table 1 in the rebuttal PDF with Table 1 in the manuscript, we can see that 4-bit quantization of eigenvector matrix $U$ has smaller quantization errors than 8-bit quantization of the preconditioner $A$. For a clearer comparison, we extract the quantization errors using Linear-2 quantization from the two tables. The results are given in the following table, where $A = A_1$ is derived from the real world as described in Subsection 3.1. We will include these results into the revision.

3) Add a description of DT quantization. Thanks. Dynamic tree (DT) quantization for $b$-bit quantization maps $\\{0, 1, \dots, 2^b - 1\\}$ to $\\{0, 1\\}\cup G$, where $G$ is a set of numbers with the following properties: the number in $G$ looks like $\pm q_k \times 10^{-E}$, where

a) $b = 2 + E + F$, where $E, F$ are integers;

b) $q_k = (p_k + p_{k+1}) / 2$, where $k \in \\{0, \dots, 2^F-1\\}$;

c) $p_j = 0.9j / 2^F + 0.1$, where $j \in \\{0, \dots, 2^F\\}$.

We will add a paragraph in Appendix C for describing quantization maps mentioned in the manuscript.

4) Add a graph from $x \in [-1,1]$ to $Q(x)$ for each quantization scheme. Per your suggestion, we have drawn the graph from $x \in [-1,1]$ to $Q(x)$ for each quantization scheme, which is similar to Figure 6 in [2]. Unfortunately, the space constraint prevents its inclusion in the one-page rebuttal PDF. We will add this graph in Appendix C.

5) Why separate $T_1$ and $T_2$ in Algorithm 3. Thanks. We follow previous second-order optimizers [3, 4], and set different values for $T_1$ and $T_2$ ($T_1$=200, $T_2$=2000 in [4]) for more efficient training. This is because the inverse root is much more expensive than preconditioner update, and thus is updated more lazily. Here we keep the default values of $T_1$ and $T_2$ for both 4- and 32-bit Shampoo for a fair comparison.

6) English improvement. Thank you for your meticulous proofreading. We will make every effort to polish our manuscript and also seek assistance from native English speakers.

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

[2] Tim Dettmers, Mike Lewis, Sam Shleifer, Luke Zettlemoyer. 8-bit optimizers via block-wise quantization. In ICLR, 2022.

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

[4] Hongwei Yong, Ying Sun, and Lei Zhang. A general regret bound of preconditioned gradient method for DNN training. In CVPR, 2023.

Thank you for the insightful and valuable comments.

1) Evaluation scope. Thanks. We trained medium-sized language models, including 124M GPT-2 and LLAMA-2 130M/350M. The results are reported in Table 2 and Figure 1 of the rebuttal PDF, which are consistent with those of vision tasks. Due to the complexity and limited rebuttal time, we were unable to conduct reinforcement learning experiments. Since our algorithm design is not task-specific, we believe its improvements should be consistently achievable on other tasks as well.

2) Time cost introduced by quantization and dequantization. Thanks. Block-wise quantization and dequantization are indeed very fast, since they can use the cuda-kernel level parallelism. As a result, they are widely adopted in activation compressed training like [1] and memory efficient optimizers like [2]. As shown in Figure 1(a), the quantization and dequantization in AdamW+4-bit Shampoo (naive) only bring around 3% extra overhead, while greatly improving memory efficiency by 24%.

3) Quantization error analysis over training. Thanks. Quantization-based approaches always try to reduce the quantization errors compared to their full-precision counterparts, so as to expect similar performance. Quantization errors will accumulate along with the training, but they are controllable. As evident from the training curves in Figures 1 and 4, no jump occurs. We experimentally analyzed the impact of quantization errors on the training process (see Appendix D2). The quantization errors of our 4-bit Shampoo are lower than those of naive Shampoo in most of cases and these errors are all bounded.

4) Time cost introduced by orthogonality rectification. Thanks. We follow the vanilla Shampoo method and perform orthogonality rectification (OR) every $T_1$ or $T_2$ iterations ($T_1$ = 100/200, $T_2$ = 500/1000). See line 5 and line 9 in Algorithm 3. This lazy update strategy greatly reduces the computational cost. We trained ResNet34 on CIFAR-100 with 4-bit Shampoo, both with and without OR, resulting in training times of 161.7 and 161.0 minutes, respectively, with OR introducing less than 1% extra overhead.

5) Comparison with Adagrad, M-FAC, and EVA. Thanks. For Adagrad, it is not as widely used as AdamW due to its inferior performance (e.g. see Table 7 in [2]). We integrate Shampoo with Adagrad, and run Adagrad for 150 epochs and Adagrad+Shampoo for 100 epochs. The table below shows that when training Swin-Tiny on CIFAR-100, Adagrad + 4 bit Shampoo converges faster than Adagrad with ignorable memory overhead, and also has higher test accuracy.

For M-FAC, the following table shows that our SGDM+4bit Shampoo enjoys much higher efficiency than M-FAC (m=32) on ResNet34, and enjoys higher test accuracy. This is because M-FAC needs to maintain m dense gradient copies (m=1024 in its official code), and is not memory-efficient. Here we run all the optimizers for 200 epochs. One can observe that even M-FAC using m=32 already requires much more GPU memory, let alone M-FAC using m=1024, and also suffers from worse performance than SGDM and SGDM+32/4 bit Shampoo.

Regarding EVA, it is a rank-one second-order optimizer and is memory-efficient. We trained ResNet34 on CIFAR-100 with SGDM+EVA, but despite extensive hyper-parameter tuning, we failed to achieve acceleration over SGDM. Instead, we cited EVA's result of training VGG-19 on CIFAR-100 for 200 epochs (see Table 2 in [5]). The test accuracies of SGDM+EVA and SGDM+Shampoo are 73% and 74.5%, respectively.

6) Theoretical analysis, e.g., convergence guarantees & mathematical justification for quantization.

Regarding convergence, we do not provide an explicit proof in this work, as our primary focus is on the practical application of optimizers to reduce GPU memory usage. This approach aligns with many other works on optimizers that prioritize practical solutions over theoretical analysis, such as [2][4].

However, we find that the proof techniques of Theorem 1 in [3] can be adapted to demonstrate the convergence of our optimizer. Specifically, we can generalize Lemma 2 used in proving Theorem 1 in [3] and then use this extended lemma to prove the convergence of quantized Shampoo. For further details, please refer to Lemma 1 in the "global” response. Note that EVA does not prove convergence, but rather demonstrates that its update step is more aggressive than that of K-FAC from the trust region optimization perspective.

In terms of quantization, we do provide a mathematical justification in Section 4 of our manuscript. We observe that the eigenvalues of Shampoo's preconditioner decrease very rapidly. So we assume that the eigenvalues follow a certain distribution, and theoretically analyze why quantizing the eigenvector matrix of the preconditioner outperforms direct preconditioner quantization.

[1] Division: memory efficient training via dual activation precision. ICML, 2023.

[2] 8-bit optimizers via block-wise quantization. ICLR, 2022.

[3] Block low-rank preconditioner with shared basis for stochastic optimization. NeurIPS, 2023.

[4] Adafactor: Adaptive learning rates with sublinear memory cost. ICML, 2018.

[5] Eva: Practical second-order optimization with Kronecker-vectorized approximation. ICLR, 2023.

Thank you for the insightful and positive comments. In the following, we provide our point-by-point response and hope our response helps address your concerns. We also look forward to the subsequent discussion which further helps to solve the current issues.

1) Breakdown of memory usage. Thanks. Per your suggestion, in the table below, we have reported the additional memory cost incurred by injecting the vanilla 32-bit Shampoo and our proposed 4-bit Shampoo into the vanilla optimizer. The extra memory cost primarily comes from the preconditioning matrices L and R, and their inverse roots. See Eqn. (1).

To compute this extra memory cost, we use the data from Table 2 of manuscript, and subtract the total memory cost (TMC) of the vanilla optimizer "A" from the TMC of "A + 32/4-bit Shampoo", where A can be SGDM and AdamW. Comparatively, the extra memory cost of the 32-bit Shampoo is seven times greater than that of our 4-bit Shampoo, demonstrating the superior efficiency of our 4-bit Shampoo. This high efficiency is also analyzed and emphasized in Appendix F.