title: 4-bit Shampoo for Memory-Efficient Network Training
summary: 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 demonstrates that our
  4-bit Shampoo achieves comparable test accuracy to its 32-bit counterpart while
  being more memory-efficient. The source code will be made available.
paper_authors._id: 665d3bf0957df09a072cd3fa
paper_authors.name: Sike Wang
paper_authors._id: 665d3bf0957df09a072cd3fb
paper_authors.name: Jia Li
paper_authors._id: 665d3bf0957df09a072cd3fc
paper_authors.name: Pan Zhou
paper_authors._id: 665d3bf0957df09a072cd3fd
paper_authors.name: Hua Huang
paper_title: 4-bit Shampoo for Memory-Efficient Network Training
paper_summary: 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 demonstrates that our
  4-bit Shampoo achieves comparable test accuracy to its 32-bit counterpart while
  being more memory-efficient. The source code will be made available.
paper_ai_keywords: second-order optimizers preconditioner first-order optimizers compressed optimizer states 4-bit optimizers 4-bit Shampoo eigenvector matrix orthogonality linear square quantization dynamic tree quantization test accuracy memory-efficient