Memory Efficient Optimizers with 4-bit States

目录

• 概
• LPMM
• 代码

Li B., Chen J. and Zhu J. Memory efficient optimizers with 4-bit states. NeurIPS, 2023.

概

本文介绍了一种支持 4-bit 的优化器量化方法.

LPMM

• 这篇文章的工作主要继承自 [8-bit optimizer].

• 本文的针对性改进主要有:

  1. 对于 First moment, 依旧采用 dynamic exponent/tree quantization, 区别是, block size 从 2048 变到 128, 从而能够规避更多的 outliers.
  2. 对于 Second moment, 用的移除 0 点的 Linear mapping. 特别的, 对于 Second moment, 进一步提出了 rank-1 normalization: 对于一个 tensor \(\mathbf{x} \in \mathbb{R}^{n \times m}\), 计算 row-wise 和 col-wise 的最大值:
     \[\mathbf{r} \in \mathbb{R}^n, \quad r_i = \max_{1 \le j \le m} x_{ij}, \\ \mathbf{c} \in \mathbb{R}^m, \quad c_j = \max_{1 \le i \le n} x_{ij}. \]
     然后通过如下方式 normalization:
     \[\frac{1}{\min\{r_i, c_j\}} x_{ij}. \]
     显然这种方式, 能够进一步缩小范围.

代码

[official-code]