SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zirui Liu $^{1*}$ , Guanchu Wang $^{1*}$ , Shaochen Zhong $^{1}$ , Zhaozhuo Xu $^{2}$ , Daochen Zha $^{1}$ , Ruixiang Tang $^{1}$ , Zhimeng Jiang $^{3}$ , Kaixiong Zhou $^{1}$ , Vipin Chaudhary $^{4}$ , Shuai Xu $^{4}$ , Xia Hu $^{1}$ $^{1}$ Rice University,  $^{2}$ Stevens Institute of Technology,  $^{3}$ Texas A&M University,  $^{4}$ Case Western Reserve University
{zl105, gw22, hv88, daochen.zha, ruixiang.tang, kaixiong.zhou, xia.hu}@rice.edu; zxu79@stevens.edu; zhimengj@amu.edu; {vipin, sxx214}@case.edu
As the model size grows rapidly, fine-tuning the large pre-trained language model has become increasingly difficult due to its extensive memory usage. Previous works usually focus on reducing the number of trainable parameters in the network. While the model parameters do contribute to memory usage, the primary memory bottleneck during training arises from storing feature maps, also known as activations, as they are crucial for gradient calculation. Notably, machine learning models are typically trained using stochastic gradient descent. We argue that in stochastic optimization, models can handle noisy gradients as long as the gradient estimator is unbiased with reasonable variance. Following this motivation, we propose a new family of unbiased estimators called WTA-CRS, for matrix production with reduced variance, which only requires storing the sub-sampled activations for calculating the gradient. Our work provides both theoretical and experimental evidence that, in the context of tuning transformers, our proposed estimators exhibit lower variance compared to existing ones. By replacing the linear operation with our approximated one in transformers, we can achieve up to  $2.7 \times$  peak memory reduction with almost no accuracy drop and enables up to  $6.4 \times$  larger batch size. Under the same hardware, WTA-CRS enables better down-streaming task performance by applying larger models and/or faster training speed with larger batch sizes. The code is available at https://github.com/zirui-ray-liu/WTACRS/.
Pre-trained language models (LMs) with transformer architecture have achieved remarkable success in numerous natural language processing (NLP) tasks [1, 2, 3, 4, 5]. Specifically, these models are trained on vast text corpora to acquire general-purpose representations, which are then adapted to a specific task by fine-tuning on task-specific data. In recent studies, it has been convincingly demonstrated that significantly increasing the number of parameters in pre-trained LMs leads to remarkable improvements in performance [6, 7]. As a result, there is now an urgent necessity to effectively adapt these models, equipped with billion-scale parameters, to a wide range of tasks.
However, a significant disparity exists between the memory requirements of pre-trained LMs and the capacity of current hardware, particularly GPUs. For example, even a GPU with 24GB memory cannot accommodate the fine-tuning process of the T5-3B model [3] with batch size one, which boasts three billion parameters. Without additional techniques, attempting to fine-tune billion-scale LMs on a single GPU is impossible. Although model-parallel fine-tuning is feasible, the majority of the time, we cannot bear the expense of acquiring multiple GPUs or the communication overhead
involved. To ensure the smooth deployment of language models during the fine-tuning process, it is crucial to adapt them for operation on a single GPU.
To address this issue, several parameter-efficient tuning methods are proposed [8, 9, 10, 11, 12, 13, 14]. Specifically, adapters [14, 13] insert a small module into the transformer blocks and only update it while keeping other parameters fixed. Similarly, prompt tuning [8] introduces a small vector that is concatenated with the input embeddings and updated during the tuning process. LoRA [12] injects trainable rank decomposition matrices into the transformer block, updating them while freezing the others. Parameter-efficient tuning methods mainly reduce the memory taken by the optimizer states [15, 12]. Although the optimizer states contribute to the memory footprint, storing activations (or feature maps) is the main memory bottleneck during training (often  $>70\%$  ) [16, 17, 18, 19]. Thus, parameter-efficient methods often do not reduce memory usage by much [9, 19].
In parallel, we can reduce the main memory bottleneck by reducing the activation storage in fine-tuning. Since transformer-based models are mainly built based on the linear layer, a less-explored direction is to replace the expensive matrix multiplication operation with its memory-efficient estimations using column-row sampling (CRS) [21, 22]. The key idea of CRS is to sub-sample tensors onto low-dimensional spaces and perform the original operations here. Specifically, for the linear operation between two matrices  $\mathbf{A} \in \mathbb{R}^{n \times m}$  and  $\mathbf{B} \in \mathbb{R}^{m \times q}$  (in the context of machine learning,  $\mathbf{A}$  is often activations), we first sample  $k$  ( $k < m$ ) column-row pairs according to a pre-defined distribution. Then we obtain  $\mathbf{A}' \in \mathbb{R}^{n \times k}$  and  $\mathbf{B}' \in \mathbb{R}^{k \times q}$  ( $k < m$ ) by picking  $k$  columns of  $\mathbf{A}$  and the corresponding rows of  $\mathbf{B}$  according to the sampled column-row pairs [22]. Finally, we estimate  $\mathbf{AB} \approx \mathbf{A}'\mathbf{B}'$ . In this way, we only need to store the sub-matrix  $\mathbf{A}'$  and  $\mathbf{B}'$  in GPU memory to perform the computation. Moreover, transformer-based models training/tuning are performed with the first-order stochastic optimizer, e.g., Adam [15]. In stochastic optimization, models can work with noisy gradients, as long as the gradient estimator is unbiased and has a reasonable variance. In view of such, we ask: why spend resources on obtaining exact gradients when we are using stochastic optimization? Motivated by this, we focus on obtaining unbiased gradients cheaply with approximated matrix multiplication.
The approximation method reduces the memory usage at the cost of giving outputs with variance. Thus there naturally exists an accuracy-memory trade-off. The main challenge is how to integrate the approximated matrix multiplication into transformer with minimal gradient variance. In this paper, we propose a new family of unbiased estimator for matrix multiplication with reduced variance, dubbed Winner-Take-All Column-Row Sampling (WTA-CRS). Compared to CRS, WTA-CRS reduces the variance of an estimator by focusing more on high-probability regions of the sampling distribution. Moreover, WTA-CRS can serve as a drop-in replacement for the linear operation in transformers, providing an unbiased weight gradient with reduced memory usage. As shown in Figure 1, our method achieves better accuracy-memory trade-off than
state-of-the-art memory-efficient tuning methods, e.g., LST [9] and LoRA [12]. Moreover, since WTA-CRS executed at the operation level, it is orthogonal to most of the existing parameter-efficient tuning methods. Our contributions are highlighted as follows:
Fig. 1. Accuracy-memory trade-off of WTA-CRS and other memory-efficient tuning methods. Unless specially stated, we use the T5-Large in the figure.
- We design a new family of unbiased estimator for matrix multiplication with reduced variance. We theoretically and experimentally verify that it has smaller variance than the established one under the context of tuning transformer.
- By replacing the linear operation with WTA-CRS in transformers, we can achieve up to  $2.7 \times$  peak memory reduction with almost no accuracy drop, and enables up to  $6.4 \times$  larger batch size. For example, to fully tune T5-3B, it requires 66.4GB of memory for in a mini-batch size of 32, which relies on A100 GPUs with a capacity of 80GB. With WTA-CRS, it only requires 21.6GB of memory for fine-tuning, which can run on a GPU with 24GB memory, e.g. RTX3090Ti.
- We implement WTA-CRS as a ready-to-use extension for Pytorch with an easy-to-use API that can also be combined with other memory-saving techniques.
In this section, we first analyze the memory usage of transformers. Then we introduce the background on the approximated matrix multiplication.
2.1 The Memory Usage of Transformers
In each training step of backpropagation, it has exactly two phases, i.e., one forward phase and one backward phase. Transformer-based models are mainly built based on the linear operation, which can be written as:
Forward Pass  $\mathbf{Z} = \mathrm{GEMM}(\mathbf{H},\mathbf{W}),$  (1a)
Backward Pass  $\nabla \mathbf{H} = \mathrm{GEMM}(\nabla \mathbf{Z},\mathbf{W}^{\top}),$  (1b)
$$
\nabla \mathbf {W} = \operatorname {G E M M} \left(\mathbf {H} ^ {\top}, \nabla \mathbf {Z}\right), \tag {1c}
$$
where  $\mathrm{GEMM}(\cdot ,\cdot)$  is the General Matrix Multiplication operation,  $\mathbf{H}$  and  $\mathbf{Z}$  are the activation (or input feature maps) and output
feature maps, respectively.  $\mathbf{W}$  is the weight of the linear layer.  $\nabla \mathbf{H}$ ,  $\nabla \mathbf{W}$ , and  $\nabla \mathbf{Z}$  are the gradient of  $\mathbf{H}$ ,  $\mathbf{W}$ , and  $\mathbf{Z}$ , respectively. From Equation (1c), activations  $\mathbf{H}$  are used in the backward phase. In commonly used deep learning framework [24, 25], it requires storing  $\mathbf{H}$  in GPU memory during the forward pass, for calculating the weight gradient  $\nabla \mathbf{W}$  in the backward pass.
Previous works show that although the model parameters contribute to the memory footprint, activations (e.g., storing  $\mathbf{H}$ ) are the main memory bottleneck during training [16, 17, 18, 19]. To get a sense of the scale, we show in Figure 2 that for popular transformer models like T5, activations may take roughly  $73\sim 88\%$  of the total memory, depending on the batch size  $B$  and sequential length  $S$ .
Fig. 2. The GPU memory usage breakdown for fine-tuning T5 [23], where the batch size  $B$  is 64 and sequential length  $S$  is 128 or 256.
2.2 Approximated GEMM With Sampling
Let  $\mathbf{X} \in \mathbb{R}^{n \times m}$ ,  $\mathbf{Y} \in \mathbb{R}^{m \times q}$  be two matrices. The goal is to efficiently estimate the matrix production  $\mathbf{XY}$ . Singular Value Decomposition (SVD) outputs provably optimal low-rank estimation of  $\mathbf{XY}$  [21]. However, SVD is almost as expensive as matrix production itself. Instead, the sampling algorithm is proposed to approximate the matrix product  $\mathbf{XY}$  by sampling  $k$  columns of  $\mathbf{X}$  and corresponding rows of  $\mathbf{Y}$  to form smaller matrices, which are then multiplied as usual [22, 26]:
$$
\operatorname {G E M M} (\mathbf {X}, \mathbf {Y}) = \sum_ {i = 1} ^ {m} \mathbf {X} _ {:, i} \mathbf {Y} _ {i,:} \approx \sum_ {t = 1} ^ {k} \frac {1}{k p _ {i _ {t}}} \mathbf {X} _ {:, i _ {t}} \mathbf {Y} _ {i _ {t,:}} = \mathbf {X} ^ {\prime} \mathbf {Y} ^ {\prime}, \tag {2}
$$
where  $\mathbf{X}_{:,i} \in \mathbb{R}^{n \times 1}$  and  $\mathbf{Y}_{i,:} \in \mathbb{R}^{1 \times q}$  are the  $i^{\text{th}}$  column and row of  $\mathbf{X}$  and  $\mathbf{Y}$ , respectively. In this paper, we call  $(\mathbf{X}_{:,i}, \mathbf{Y}_{i,:})$  the  $i^{\text{th}}$  column-row pair.  $k$  is the number of sampled pairs ( $1 \leq k \leq m$ ).  $\mathcal{P} = \{p_i\}_{i=1}^m$  is a probability distribution over the column-row pairs.  $i_t \in \{1, \dots, m\}$  is the index of the sampled column-row pair at the  $t^{\text{th}}$  trial.  $s_t$  is the scale factor.  $\mathbf{X}' \in \mathbb{R}^{n \times k}$  and  $\mathbf{Y}' \in \mathbb{R}^{k \times q}$  are the normalized sub-matrices sliced according to the sampled column-row pairs.
Existing work [22] shows  $\mathbf{X}'\mathbf{Y}'$  is an unbiased estimation of  $\mathbf{X}\mathbf{Y}$ , i.e.,  $\mathbb{E}[\mathbf{X}'\mathbf{Y}'] = \mathbf{X}\mathbf{Y}$ . Furthermore, the approximation error  $\mathbb{E}[||\mathbf{X}\mathbf{Y} - \mathbf{X}'\mathbf{Y}'||_F]$  is minimized when the probabilities  $\{p_i\}_{i=1}^m$  are proportional to the product of the column-row Euclidean norms [22] (Proof in Appendix C):
$$
p _ {i} = \frac {\left| \left| \mathbf {X} _ {: , i} \right| \right| _ {2} \left| \left| \mathbf {Y} _ {i , :} \right| \right| _ {2}}{\sum_ {j = 1} ^ {m} \left| \left| \mathbf {X} _ {: , j} \right| \right| _ {2} \left| \left| \mathbf {Y} _ {j , :} \right| \right| _ {2}}. \tag {3}
$$
As we analyzed in Section 2.1, storing the activation  $\mathbf{H}$  is the major memory bottleneck. If we can replace  $\mathrm{GEMM}(\mathbf{H}^{\top},\nabla \mathbf{Z})$  in Equation (1c) with  $\mathbf{H}'^\top \nabla \mathbf{Z}'$  following the paradigm of Equation (2), then we only need  $\mathbf{H}'$  instead of  $\mathbf{H}$  in GPU memory to compute the gradient, which significantly decreases the memory usage of activations. This estimation linearly reduces the memory complexity from  $\mathcal{O}(nm)$  to  $\mathcal{O}(nk)$ . Also, the total number of floating point operations (FLOPs) is reduced as
well since the computation is executed on two smaller matrices. For the ease of illustration, in this paper we call the distribution in Equation (3) the column-row index distribution. In the next section, we explore how to reduce memory usage via sampling-based matrix multiplication.
In recent years, we have observed that deep neural network training can be performed almost entirely with first-order stochastic optimization [15]. Thus intuitively, in stochastic optimization we can reduce the resources spent on obtaining gradients, as long as the estimated gradient is unbiased with reasonable variance [27, 28, 29]. Following this motivation, we first design a new unbiased estimator for matrix multiplication with reduced variance compared to the one in Equation (2) (Section 3.1). Then we introduce how to replace the GEMM in Transformer with its approximated version to reduce the memory usage (Section 3.2).
3.1 Winner-Take-All Column-Row Sampling: A New Unbiased Estimator for GEMM
In this section, we mathematically design a new unbiased estimator for GEMM with reduced variance called WTA-CRS (Winner-Take-All Column-Row Sampling). Following the notation in Section 2.2, let  $\mathbf{X} \in \mathbb{R}^{n \times m}$ ,  $\mathbf{Y} \in \mathbb{R}^{m \times q}$  be two matrices.  $\mathcal{P} = \{p_i\}_{i=1}^m$  is the column-row index distribution in Equation (3) $^2$ . We first define the variable  $f(i)$  as  $f(i) = \frac{\mathbf{X}_{:i}\mathbf{Y}_{:i}}{p_i}$ .
$f(i)$  is an unbiased estimation for the matrix production between  $\mathbf{X}$  and  $\mathbf{Y}$ . To see this,
$$
\mathbb {E} _ {j \sim \mathcal {P}} [ f (j) ] = \sum_ {i = 1} ^ {m} p _ {i} \frac {\mathbf {X} _ {: , i} \mathbf {Y} _ {i :}}{p _ {i}} = \mathbf {X Y}.
$$
We note that the prior approximated matrix multiplication in Equation (2) is the direct extension of  $f(i)$  by taking the average of  $\{f(i_t)\}_{t=1}^k$  among  $k$  independent random trials to reduce the variance. Here we explore an alternative approach to reduce the variance of  $f(i)$  beyond simple averaging. Our core idea is to partition the column-row index distribution  $\mathcal{P} = \{p_i\}_{i=1}^m$  into two complementary regions based on the probability mass: a high-probability region  $\mathcal{P}^C$  and a low-probability region  $\mathcal{P}^{D \setminus C}$ , where  $\mathcal{D} = \{1, \dots, m\}$  is the whole set and  $C$  is the set of the column-row index with the largest probability. Let  $C$  be the set of column-row pair indices associated with  $|C|$  largest  $p_i$ . We define WTA-CRS estimator for  $\mathbf{X}\mathbf{Y}$  as follows:
$$
\mathbb {E} _ {j \sim \mathcal {P D} \backslash c} \left[ \sum_ {c \in \mathcal {C}} f (c) p _ {c} + \left(1 - \sum_ {c \in \mathcal {C}} p _ {c}\right) f (j) \right]. \tag {4}
$$
We note that the random variable in Equation (4) is the column-row pair index  $j$ , and is only sampled from  $\mathcal{D} \backslash \mathcal{C}$ . The estimator defined in Equation (4) contains two parts. The first part  $\sum_{c \in \mathcal{C}} f(c)p_c$  has no relationship with the random variable  $j$  and is summed deterministically. The second part  $f(j)$  is sampled stochastically, but scaled by the factor  $(1 - \sum_{c \in \mathcal{C}} p_c)$ . When  $\mathcal{P} = \{p_i\}_{i=1}^m$  is concentrated on a small number of atoms, the scaling factor  $(1 - \sum_{c \in \mathcal{C}} p_c)$  for the stochastic term should be small. Therefore, we intuitively expect the estimator to have a small variance in this case due to a small scaling factor. In this way, we reduce the variance of an estimator by focusing more on high-probability regions of the distribution (winner-take-all). Below we formalize this intuition by showing the statistical property of our estimator regarding the bias and variance, respectively.
Theorem 1 (Proof in Appendix C.2). The estimator defined in Equation (4) is an unbiased estimator for matrix production  $\mathbf{X}\mathbf{Y}$ , i.e.,  $\mathbb{E}_{j\sim \mathcal{P}^{\mathcal{D}}\backslash c}[\sum_{c\in \mathcal{C}}f(c)p_c + (1 - \sum_{c\in \mathcal{C}}p_c)f(j)] = \mathbf{X}\mathbf{Y}$ .
Theorem 1 states that our proposed estimator in Equation (4) is unbiased. Below we compare our proposed estimator to the CRS estimator in Equation (2) in terms of the variance. Suppose we have the budget of only utilizing  $k$  column-row pairs for approximating the matrix production. From the implementation perspective, the estimator defined in Equation (2) estimates  $\mathrm{GEMM}(\mathbf{X},\mathbf{Y})$  as:
$$
\left(\text {C R S}\right) \quad g (\mathbf {X}, \mathbf {Y}) = \frac {1}{k} \sum_ {t = 1} ^ {k} f \left(i _ {t}\right), \quad i _ {1}, \dots , i _ {k} \stackrel {\text {i . i . d}} {\sim} \mathcal {P}. \tag {5}
$$
Here we note that the theoretical analysis in this section can be applied to any probability distribution, not only limited to the one in Equation (3).
Our estimator defined in Equation (4) splits the budget  $k$  into two parts. Namely, the first part explicitly sums the expectation terms for the largest probability group  $\mathcal{C}(|\mathcal{C}| < k)$ , while stochastically average  $k - |\mathcal{C}|$  samples drawn from  $\mathcal{D} \backslash \mathcal{C}$  to estimate the remaining terms, up to scale:
$$
\left(\mathrm {W T A} - \mathrm {C R S}\right) \hat {g} (\mathbf {X}, \mathbf {Y}) = \sum_ {c \in \mathcal {C}} f (c) p (c) + \frac {1 - \sum_ {c \in \mathcal {C}} p _ {c}}{k - | \mathcal {C} |} \sum_ {j = 1} ^ {k - | \mathcal {C} |} f (j), \quad i _ {1}, \dots , i _ {k - | \mathcal {C} |} \stackrel {\text {i . i . d}} {\sim} \mathcal {P} ^ {\mathcal {D} \backslash \mathcal {C}}. \tag {6}
$$
Theorem 2 (Proof in Appendix C.3). Suppose the total budget of column-row pairs is  $k$ . If  $\mathcal{C}$  satisfies
$$
\sum_ {c \in \mathcal {C}} p _ {c} > \frac {| \mathcal {C} |}{k}, \tag {7}
$$
then we have  $\operatorname{Var}[\hat{g}(\mathbf{X},\mathbf{Y})] < \operatorname{Var}[g(\mathbf{X},\mathbf{Y})]$ . Moreover,  $\operatorname{Var}[\hat{g}(\mathbf{X},\mathbf{Y})]$  is minimized when  $|\mathcal{C}| = \min_{|\mathcal{C}| \in \{0,\dots,k\}} \frac{1 - \sum_{c \in \mathcal{C}} p_c}{k - |\mathcal{C}|}$ .
Both the left- and right-hand sides of Equation (7) depend on the size of the highest probability group  $|\mathcal{C}|$ , which controls the number of high probability column-row pairs that are directly added without sampling. Below we experimentally investigate whether Equation (7) holds under the context of fine-tuning the transformer-based model with varying  $|\mathcal{C}|$ .
Fig. 3. The probability mass  $\sum_{c\in \mathcal{C}}p_c$  versus  $\frac{|C|}{k}$  in Equation (7) at  $k = 0.3|\mathcal{D}|$ . Here we visualize the column-row index distribution of query/key/value projection layer in the T5-base model, which is fine-tuned on RTE dataset. More similar results can be found in Appendix E.1.
Experimental analysis. As shown in Figure 3, we visualize the two terms in Equation (3) for the column-row index distribution of query, key, and value projection in the self-attention module, respectively [1]. Specifically, we fix the total column-row pair budget  $k = 0.3|\mathcal{D}|$  and change the size of the highest probability group  $|\mathcal{C}|$  from 0 to  $k$ . We conclude that Equation (7) holds for most of the layers when fine-tuning transformers. Thus, we expect our WTA-CRS has better performance than CRS for adapting transformer-based models, which is later experimentally verified in Section 5.
3.2 Compress GEMM in Transformers with WTA-CRS
Previous work has shown that unbiasedness of the estimated gradient is crucial for the proper convergence of stochastic gradient descent [27, 28, 16, 30]. As shown in Section 2.1, we have three GEMM in the linear layer. Below we investigate how to replace GEMM with its approximated version in a way that the estimated gradient is unbiased.
Unbiasedness. Previous work has shown that to ensure the unbiasedness of the gradient, the approximation can only be applied during the backward pass [16, 31, 21]. The rationale behind this conclusion is that we have  $\mathbb{E}[f(x)]\neq f(\mathbb{E}[x])$  for any non-linear function  $f(\cdot)$ , e.g.,
$\mathbb{E}[x^2 ]\neq \mathbb{E}^2 [x]$ . Thus if we replace the forward GEMM in Equation (1a), even when the approximation method gives an unbiased estimation, i.e.,  $\mathbb{E}[\hat{g} (\mathbf{H},\mathbf{W})] = \mathbf{H}\mathbf{W} = \mathbf{Z}$ , the output activations (e.g.,  $\mathrm{GeLU}(\mathbf{Z}))$  are still biased since the activation function is non-linear, namely,
$$
\operatorname {G e L U} (\hat {g} (\mathbf {H}, \mathbf {W})) = \operatorname {G e L U} (\mathbb {E} [ \mathbf {Z} ]) \neq \mathbb {E} [ \operatorname {G e L U} (\mathbf {Z}) ].
$$
Fig. 4. The illustration of how to deploy WTA-CRS to linear layers. We only replace GEMM in Equation (1c) with its approximated version using WTA-CRS. The pseudocode is given in Appendix D Algorithm 1.
Fig. 5. The diagram of a single Transformer block. The shape of activations is annotated, where  $B, S, D_{\mathrm{model}}$ ,  $N_{\mathrm{head}}$ , and  $D_{\mathrm{head}}$  are the batch size, sequence length, hidden size, number of attention heads, and head dimension, respectively. WTA-CRS can be applied to the operators in green; the activation maps of operators in blue can be losslessly compressed; and those in gray are not compressed in this paper. The idea of this figure is inspired by [32].
To ensure the unbiasedness of gradient and reduce the memory usage of storing  $\mathbf{H}$ , as shown in the example of Figure 4, we only replace GEMM in the backward pass with its approximation (e.g., Equation (1c)), while leaving the forward one unchanged (e.g., Equation (1a)). We show in Appendix B that the estimated weight gradient is unbiased in this case.