SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Bingrui Li $^{1}$ , Jianfei Chen $^{1\dagger}$ , Jun Zhu $^{1}$
$^{1}$ Dept. of Comp. Sci. and Tech., Institute for AI, BNRist Center, THBI Lab, Tsinghua-Bosch Joint ML Center, Tsinghua University
lbr22@mails.tsinghua.edu.cn; {jianfeic, dcszj}@tsinghua.edu.cn
Optimizer states are a major source of memory consumption for training neural networks, limiting the maximum trainable model within given memory budget. Compressing the optimizer states from 32-bit floating points to lower bitwidth is promising to reduce the training memory footprint, while the current lowest achievable bitwidth is 8-bit. In this work, we push optimizer states bitwidth down to 4-bit through a detailed empirical analysis of first and second moments. Specifically, we find that moments have complicated outlier patterns, that current block-wise quantization cannot accurately approximate. We use a smaller block size and propose to utilize both row-wise and column-wise information for better quantization. We further identify a zero point problem of quantizing the second moment, and solve this problem with a linear quantizer that excludes the zero point. Our 4-bit optimizers are evaluated on a wide variety of benchmarks including natural language understanding, machine translation, image classification, and instruction tuning. On all the tasks our optimizers can achieve comparable accuracy with their full-precision counterparts, while enjoying better memory efficiency.*
Large-scale models with a massive amount of parameters [5, 9, 20, 22, 49, 58] have shown impressive few-shot learning abilities on general tasks [52]. Despite being powerful, training these models is challenging. Memory capacity is one of the main bottlenecks of training large-scale models. Modern neural networks are typically trained with stateful optimizers such as Adam [26], which need to maintain one or two optimizer states (i.e., first and second moments) per each parameter. As the model size grows, the memory consumed by optimizer states can be a dominating factor of memory consumption [40, 41].
There are several attempts to reduce optimizers' memory consumption. Factorization [2, 8, 46] applies low-rank approximation to optimizer states, delta tuning [18, 23, 24, 27, 28] avoids maintaining optimizer states for most parameters by only tuning a small subset, and low-precision optimizers [15, 44] represent their states with low-precision numerical formats, which consume less memory.
Among these methods, low-precision optimizers are attractive due to their simplicity and wide applicability. Dettmers et al. [15] propose an 8-bit optimizer with reparameterized embedding layers ("stable embedding layers") and a block-wise 8-bit dynamic exponential numerical format for optimizer states. Their 8-bit optimizers achieve similar convergence to full-precision optimizers on language modeling, image classification, machine translation, and language understanding tasks.
In this work, we push the required numerical precision of low-precision optimizers from 8 to 4-bit through analyzing the patterns in the first and second moments and designing dedicated quantizers
*Code is available at https://github.com/thu-ml/low-bit-optimizers
Algorithm 1 Compression-based Memory Efficient Optimization Framework
Require: black-box stochastic optimization algorithm  $\mathcal{A}$ , initial parameter  $\mathbf{w}_0\in \mathbb{R}^p$ , initial optimizer state  $\bar{\mathbf{s}}_0 = \mathbf{0}$ , total number of iterations  $T$   
1: for  $t = 1,2,\ldots ,T$  do  
2: Sample a minibatch  $\zeta_t$  and get stochastic gradient  $\mathbf{g}_t = \nabla_{\mathbf{w}}f(\mathbf{w}_{t - 1},\zeta_t)$   
3:  $\mathbf{s}_{t - 1}\gets \mathrm{decompress}(\bar{\mathbf{s}}_{t - 1})$   
4:  $\mathbf{w}_t,\mathbf{s}_t\gets \mathcal{A}(\mathbf{w}_{t - 1},\mathbf{s}_{t - 1},\mathbf{g}_t)$   
5:  $\bar{\mathbf{s}}_t\gets \mathrm{compress}(\mathbf{s}_t)$   
6: end for  
7: return  $\mathbf{W}_T$
for optimizer states. Specifically, we find that moments exhibit complicated outlier patterns, that vary across different parameter tensors. The large block size proposed by Dettmers et al. [15] cannot properly handle all different outlier patterns. Based on this observation, we propose to use a smaller block size, which improves the approximation of first moment.
For the second moment, we find its quantization suffers from a zero-point problem. Since the update direction is usually inversely proportional to the square root of the second moment, quantizing non-zero quantities to zero will cause significant deviation. To address this problem, we propose a simple linear quantizer to exclude the zero-point for second moment. We further propose a stronger quantization technique, rank-1 normalization, to improve the approximation of second moment by better handling the outlier patterns. Our proposed quantization techniques are robust enough to achieve lossless convergence under 4-bit, even without the stable embedding layers proposed by Dettmers et al. [15].
Finally, we investigate the combination of factorization methods [2, 46] with low-precision optimizers, and propose a memory efficient optimizer which utilizes quantized 4-bit first moment and factorized second moment. For applicable tasks, the hybrid optimizer enjoys best of both worlds: good convergence and memory efficiency.
We evaluate our 4-bit optimizers on a wide range of tasks, including natural language understanding, machine translation, image classification, and instruction tuning of large language models. On all the benchmarks, our 4-bit optimizer can converge similarly fast with full-precision optimizers, and the converged models do not have noticeable accuracy loss. Our optimizers consume less memory than existing 8-bit optimizers [15], while improves the throughput of language model finetuning with optimizer offloading [41, 45] due to reduced communication cost.
In this section, we present some preliminaries of compression-based memory efficient optimizers and discuss quantization methods for compression of optimizer states in a general formulation.
2.1 A Framework for Compression-based Memory Efficient Optimizers
Gradient-based stateful optimizers like SGDM [38, 47], Adam [26], AdamW [32] are the principal choices in deep neural network training. However, the memory footprint of stateful optimizers is several times of model itself, which results in a bottleneck for large model pretraining/finetuning. Consider the update rule of the Adam optimizer:
$$
\mathbf {A d a m} \left(\mathbf {w} _ {t - 1}, \mathbf {m} _ {t - 1}, \mathbf {v} _ {t - 1}, \mathbf {g} _ {t}\right) = \left\{ \begin{array}{l l} \mathbf {m} _ {t} & = \beta_ {1} \mathbf {m} _ {t - 1} + (1 - \beta_ {1}) \mathbf {g} _ {t} \\ \mathbf {v} _ {t} & = \beta_ {2} \mathbf {v} _ {t - 1} + (1 - \beta_ {2}) \mathbf {g} _ {t} ^ {2} \\ \hat {\mathbf {m}} _ {t} & = \mathbf {m} _ {t} / \left(1 - \beta_ {1} ^ {t}\right) \\ \hat {\mathbf {v}} _ {t} & = \mathbf {v} _ {t} / \left(1 - \beta_ {2} ^ {t}\right) \\ \mathbf {w} _ {t} & = \mathbf {w} _ {t - 1} - \eta \cdot \hat {\mathbf {m}} _ {t} / \left(\sqrt {\hat {\mathbf {v}} _ {t}} + \epsilon\right) \end{array} \right. \tag {1}
$$
During the training process, the model parameters  $\mathbf{w}_t$  and optimizer states (i.e., first and second moments)  $\mathbf{m}_t$ ,  $\mathbf{v}_t$  need to be stored persistently in the GPU memory. As the model grows large, optimizer states are a main source of training memory consumption. Each parameter dimension takes
two 32-bit optimizer states, making the memory consumption of Adam-style optimizers three times larger than stateless optimizers like SGD.
Compressing the optimizer states is a promising method to achieve memory efficient optimization. Formally, given a gradient-based optimizer  $\mathcal{A}$ , a memory efficient version of the optimization algorithm with compressed optimizer states is given by Alg. 1. The algorithm  $\mathcal{A}$  can be any gradient-based optimizers like SGDM, Adam, AdamW, etc. See App. F for examples. In Alg. 1, the precise states  $\mathbf{s}_t$  are only temporal, and only compressed states  $\overline{\mathbf{s}}_t$  are stored persistently in the GPU memory. Memory footprint reduction can be achieved since in neural networks, the state vectors  $\mathbf{m}_t$  and  $\mathbf{v}_t$  are usually concatenation of state vectors of each parameterized layer. Therefore, we can perform the optimizer steps (Line 3-5) separately for each layer, so only one single layer of the precise states are presented in the memory at a time. The states for all other layers are kept compressed. In the rest of this paper, we focus on the compression and decompression method, to achieve high compression rate of optimizer states while maintaining good convergence of the optimizer.
2.2 Main Compression Method: Quantization
Quantizing the optimizer states to lower precision (e.g. 8-bit integers) is an effective way to compression optimizer states [15]. In this case, the optimizer states are compressed with a quantizer and decompressed with a dequantizer. The low-precision numerical format significantly impacts the accuracy of quantization methods. Here, we present a general framework for numerical formats we considered for compressing optimizer states.
A quantizer converts full-precision tensors to low-precision formats. Based on the formulation proposed by Dettmers and Zettlemoyer [16], we disentangle the quantizer  $\mathbf{Q}(\cdot)$  into two parts: normalization  $\mathbf{N}(\cdot)$  and mapping  $\mathbf{M}(\cdot)$ , which applies sequentially and element-wisely to a tensor to be quantized. For concise presentation, we only discuss quantizers for unsigned inputs, and defer the discussion of signed inputs in App. E.1. Formally, the quantizer for a tensor  $\mathbf{x} \in \mathbb{R}^p$  is given by
$$
q _ {j} := \mathbf {Q} (x _ {j}) = \mathbf {M} \circ \mathbf {N} (x _ {j}).
$$
Normalization The normalization operator  $\mathbf{N}$  scales each element of  $\mathbf{x}$  into the unit interval, i.e. [0, 1]. Normalization can have different granularity, such as per-tensor, per-token (row) [37, 55], per-channel (column) [4], group-wise [37, 55] and block-wise [15]. The per-tensor and block-wise normalization operators are given by
$$
\begin{array}{l} n _ {j} := \mathrm {N} _ {\text {p e r - t e n s o r}} (x _ {j}) = x _ {j} / \max  _ {1 \leq i \leq p} | x _ {i} |, \\ n _ {j} := \mathrm {N} _ {\text {b l o c k - w i s e}} (x _ {j}) = x _ {j} / \max  \left\{\left| x _ {i} \right|: 1 + B \left\lfloor j / B \right\rfloor \leq i \leq B \left(\left\lfloor j / B \right\rfloor + 1\right) \right\}, \\ \end{array}
$$
respectively, where the involved scaling factors are called quantization scale, which are persistently stored together with quantized tensor until dequantization. The granularity of normalization presents a trade-off of quantization error and memory overhead. Normalization method with low quantization error and acceptable memory overhead is preferred. In this case, the coarsest per-tensor normalization operator has negligible memory overhead, i.e. only 1 scaling factor regardless of tensor size, but suffers from largest error due to outliers. Block-wise normalization views the tensor as an 1-dimensional array, divides the array into blocks of size  $B$  called block and assigns a quantization scale within each block, which leads to  $\lceil p / B \rceil$  quantization scales totally. The block size could be adapted to trade-off quantization error and memory overhead.
Mapping A mapping [15] converts normalized quantities to low-bitwidth integers. Formally, the mapping operator  $\mathbf{M} = \mathbf{M}_{\mathbf{T},b}$  is equipped with a bitwidth  $b$  and a predefined increasing mapping, named quantization mapping  $\mathbf{T}:[0,2^b -1]\cap \mathbb{Z}\to [0,1]$ . Then  $\mathbf{M}$  is defined as
$$
q _ {j} := \mathbf {M} (n _ {j}) = \arg \min  _ {0 \leq i <   2 ^ {b}} | n _ {j} - \mathbf {T} (i) |.
$$
The design of  $\mathbf{T}$  is critical as it could effectively mitigate quantization error by capturing the distribution information of  $\mathbf{n}$ . There are two kinds mappings that are of specific interest to optimizer states quantization, linear mapping and dynamic exponent (DE) mapping [13]. A linear mapping  $\mathbf{T}(i) = (i + 1) / 2^b$  defines a linear quantizer, where the quantization intervals distribute evenly in each block. The DE mapping can approximate small values well, similar to floating point numbers. DE splits the binary representation of a low-precision integer  $i$  into three parts: a leading
Figure 1: Visualization of the first moment in the layers.3blocks.1.mlp.fc1 layer in a Swin-T model. (a): Magnitude of the first moment. (b): Histogram of the first moment. (c): Moment approximated by B128/DE. (d): Moment approximated by B2048/DE.
sequence of  $E$  zeros, followed by an indicator bit one, and remaining  $F$  fraction bits. DE defines  $\mathbf{T}(i) = 10^{-E(i)}\mathrm{fraction}(i)$ , where  $\mathrm{fraction}(i)\in [0.1,1]$ . See App. E.2 for full specifications and visualizations of different quantization mappings.
The normalization  $\mathbf{N}$  and mapping  $\mathbf{M}$  roughly play the same role in finding good quantization point candidates and they affect each other. If an oracle normalization scaling the original tensor  $\mathbf{x}$  to a uniform distribution is accessible, linear mapping could be used generally. On the contrary, if the optimal mapping could be readily identified with respect to certain metrics for a per-tensor normalized tensor, there is no necessity to develop additional normalization methods that incur extra overhead for mitigating the negative effects of outliers. In essence, if one of the two operators approaches optimality, the quantizer can still perform well even when the other operator is set to its most basic configuration. Additionally, as one operator becomes increasingly powerful, the potential for improvement in the other operator gradually diminishes.
Dequantization The dequantizer is just the inverse of the quantizer, which is simply
$$
\tilde {x} _ {j} := \mathbf {N} ^ {- 1} \circ \mathbf {T} (q _ {j}).
$$
Based on our formulation, we name quantizers by their normalization and mapping methods as Norm. /Map.. For example, 8-bit optimizers [15] use block-wise normalization with a block size of 2048 and dynamic exponent mapping, which we call B2048/DE in the rest of the paper.
3 Compressing First Moment
In the next two sections, we describe our design of compression and decompression methods in Alg. 1 to realize our memory efficient 4-bit optimizers. We first discuss the compression of the first moment. Our compression method for first moment is based on Dettmers et al. [15], which uses block-wise normalization with a block size of 2048 and dynamic exponent mapping [13]. We preliminary reduce the bitwidth from 8-bit to 4-bit and discover that the first moment is rather robust to quantization. This simple optimizer can already converge with 4-bit first moment, though in some cases there is accuracy degradation.
To further improve the performance, we investigate the patterns of the first moment. There are some outliers in the moments and outliers significantly affect quantization scale due to their large magnitude. Outlier patterns have been studied for weights and activations. It is shown that the weights are rather smooth [14, 54, 57], while the activation has column-wise outliers [4, 14, 53, 54], i.e., the outliers in activation always lie in some fixed channels. However, we find that the outlier patterns in optimizer states are quite complicated. Specifically, Fig. 2 shows patterns of first moment in a Swin-T model during training. Outliers of the layers .3. blocks .0.mlp.fc1 layer lie in fixed rows while outliers of the layers .3. blocks .0.mlp.fc2 layer lie in fixed columns. Actually, the outlier patterns vary across different architectures, layers and parameters. See more patterns in App. B.
The complicated outlier pattern makes optimizer states harder to quantize. There exist some moment tensors, that the outliers persist roughly in certain columns (dimension 1), as shown in Fig. 1. Block-wise normalization treats this tensor as an flattened one-dimensional sequence, in the row-first order. Therefore, any block size of 2048 would include an entry in the outlier column, resulting
(a)
Figure 2: Outlier patterns vary across two first moment tensors. (a): outliers lie in fixed rows (dimension 0). (b): outliers lie in fixed columns (dimension 1).
Figure 3: Histogram of the inverse square root of second moment. (a) full-precision; (b) quantized with B128/DE; (c) quantized with B128/DE-0. All figures are at log10 scale and y-axis represents density.
in a large quantization scale. In this case, block-wise normalization is not better than per-tensor normalization. Therefore, we adopt a smaller block size of 128, as it provides enhanced performance while incurring only a little memory overhead. In Fig. 1, we show that quantizing with a smaller block size approximates the moments better than a large block size.
4 Compressing Second Moment
Compared to the first moment, quantizing the second moment is more difficult and incurs training instability. Besides its sharper outlier pattern and ill-conditioned distribution compared with first moment, we further identify the zero-point problem as the main bottleneck of quantizing the second moment. We propose an improved normalization method with a quantization mapping excluding zero. We also propose a factorization method for compressing the second moment, which leads to even better memory efficiency.
4.1 Zero-point Problem
The problem of quantizing second moment different from quantizing other tensors in neural networks. For weights [25], activations [25], and gradients [1], it is desirable to contain zero in the quantization mapping for lower approximation error and training stability. Empirically, zero is often the most frequent element [57]. But for second moment in Adam, small values around zero significantly impact the update direction, which is proportional to the inverse square root of the second moment, as shown in Eq. 1. In this case, a small quantization error of values around zero will cause catastrophically large deviation of the update direction. Fig. 3 shows the histogram of the inverse square root (i.e., transformed with  $h(v) = 1 / (\sqrt{v} + 10^{-6})$ ) of the second moment quantized with B128/DE. The quantizer pushes most of entries of the tensor to zero, so the inverse square root of most points fall into  $10^{6}$  due to zero-point problem, resulting in a complete degradation in approximation. One remedy is to simply remove zero from the DE quantization map, which we call DE-0. The smallest number representable by DE-0 is 0.0033. With the B2048/DE-0 quantizer applied to second moment, the approximation of second moment is more precise (Fig. 3), and the training with 4-bit optimizer stabilizes (Tab. 1). However, by removing zero, DE-0 wastes one of the  $2^{4} = 16$  quantization points. We propose to use a Linear mapping  $\mathbf{T}(i) = (i + 1) / 2^b$ , whose smallest representable number is 0.0625. The linear mapping performs better than DE-0 for quantizing second moment.
In Tab. 1, we ablate different quantization schemes and show that excluding zero from the mapping is indeed the key factor for second moment quantization, which cannot be replaced by a smaller block size or/and stochastic rounding [11]. We also test Stable Embedding layers proposed by Dettmers et al. [15], which are reparameterized embedding layers which can be more stably optimized. While Stable Embedding could improve training stability, it cannot fully retain accuracy since non-embedding layers still suffer from zero points. On the other hand, when the zero-point problem is properly addressed, Stable Embedding is no longer required to retain accuracy.
4.2 Rank-1 Normalization
We propose an empirically stronger normalization method named rank-1 normalization based on the observation of the heterogeneous outlier pattern in Sec. 3, inspired by the SM3 optimizer [2]. Formally, for a matrix-shaped non-negative (second moment) tensor  $\mathbf{x} \in \mathbb{R}^{n \times m}$ , its 1-dimensional
Table 1: Ablation analysis of 4-bit optimizers on the second moment on the GPT-2 Medium E2E-NLG finetuning task. The first line barely turns 8-bit Adam [15] into 4-bit, i.e. B2048/DE for both first and second moments. We only vary the quantization scheme for second moment. SR=stochastic rounding (see App. E.3 for details). Stable Embedding layers are not quantized. 32-bit AdamW achieves a BLEU of 67.7.
Normalization	Mapping	Stable Embed.*	Factorized	Unstable(%)	BLEU
B2048	DE	X	X	33	66.6 ± 0.61
B2048	DE	✓	X	0	66.9 ± 0.52
B128	DE	X	X	66	65.7 ± N/A
B128	DE+SR*	X	X	33	65.4 ± 0.02
B128	DE	✓	X	0	67.2 ± 1.13
B2048	DE-0	X	X	0	67.5 ± 0.97
B2048	DE-0	✓	X	0	67.1 ± 1.02
B128	DE-0	X	X	0	67.4 ± 0.59
Rank-1	DE-0	X	X	0	67.5 ± 0.58
Rank-1	Linear	X	X	0	67.8 ± 0.51
Rank-1	Linear	X	✓	0	67.6 ± 0.33
statistics  $\mathbf{r} \in \mathbb{R}^n$  and  $\mathbf{c} \in \mathbb{R}^m$  are defined as  $r_i = \max_{1 \leq j \leq m} x_{i,j}$  and  $c_j = \max_{1 \leq i \leq n} x_{i,j}$ , which are exactly the quantization scales of per-row and per-column normalization. Rank-1 normalization utilizes two scales jointly and produces a tighter bound for entry, which is defined as
$$
\mathrm {N} _ {\text {r a n k - 1}} \left(x _ {i, j}\right) = \frac {1}{\operatorname* {m i n} \left\{r _ {i} , c _ {j} \right\}} x _ {i, j}.
$$
Rank-1 normalization could be easily generalized to high-dimensional tensors and signed tensors. See App. G for details and the pseudocode.
Compared with per-tensor, per-token (row), and per-channel (column) normalization, rank-1 normalization utilizes the 1-dimensional information in a more fine-grained manner and gives element-wisely unique quantizaion scales. It deals with the outliers more smartly and effectively when outlier persists in fixed rows or columns but the pattern across tensors are unknown and/or varied (Fig. 2). On the other hand, block-wise normalization is also capable to capture the local information and avoid outliers effectively regardless of the patterns when a small block size is taken, but rank-1 normalization provides a better trade-off between memory overhead and approximation. Rank-1 normalization falls back to per-tensor normalization for 1-dimensional tensors, so we use B128 normalization in those cases. Empirically, rank-1 normalization match or exceed the performance of B128 normalization at a moderate model scale (hidden_size=1024). It is also possible to combine block-wise and rank-1 normalization together, which we leave for future work.
4.3 Factorization
Many memory efficient optimization methods [2, 8, 46] represent the entire second moment with a few number of statistics, which differ from quantization and they take only sublinear memory cost. These works bring more memory saving but they are only applicable to the second moment. In this work, we find the factorization method proposed in Adafactor [46] could also avoid zero-point problem effectively, as shown in Tab. 1. Further, we explore the effects of factorization on second moment based on quantized optimizers to attain maximal memory saving while maintain lossless accuracy. Specifically, when factorization is enabled, we factorize all second moment with dimension greater than 1, and quantize 1d second moment tensors.
We compare our 4-bit optimizers with their full-precision counterparts, as well as other memory efficient optimizers including 8-bit AdamW [15]†, Adafactor [46] and SM3 [2]. 8-bit AdamW's optimizer states are not quantized for embedding layers. For Adafactor, we compare both the
$\dagger$  https://github.com/TimDettmers/bitsandbytes