Winner-Take-All Column Row Sampling for Memory Efficient Adaptation of Language Model
Authors: Zirui Liu
Guanchu Wang
Shaochen Zhong
Zhaozhuo Xu
Daochen Zha
Ruixiang Tang
Zhimeng Jiang
Kaixiong Zhou
Vipin Chaudhary
Shuai Xu
Xia Hu
Venue: NeurIPS 2023 poster (2023)
Keywords: memory-efficient tuning
language model
transformers

Abstract: 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 \sas, 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.7X peak memory reduction with almost no accuracy drop and enables up to $6.4\times$ larger batch size.
Under the same hardware, \sas 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://anonymous.4open.science/r/WTACRS-A5C5/.

Winner-Take-All Column Row Sampling for Memory Efficient Adaptation of Language Model
Zirui Liu't, Guanchu Wang"", Shaochen Zhong', Zhaozhuo Xu², Daochen Zha', Ruixiang Tang'. Zhimeng Jiang Kaixiong Zhou', Vipin Chaudhary's Shuai Xu', Xia Hu
'Rice University, Steevens Institute of Technology. Teexas A&M University,
*Case Western Reserve University
g222, dacchen zha, ruixiang tang, kaixiong 2hou, xia. hur@riceeeeeeeee edu; zxu79@stevens. edu; zhimengj @tamu. edu; {vipin, sxx214]0case. edu
Abstract
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 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 activa- tions, 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 handle noisy gradients long the gradient estimator is unbiased with reasonable variance. Following this motivation, we propose a new family of unbiased estimators called WTA CRSS 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, the context of tuning transformers, our proposed estimators exhibit lower variance compared to existing ones. By replacing the linear operation with our approxi- mated one in transformers, we can achieve up 2.7x peak memory reduction with almost no accuracy drop and enables up to 6.4 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 hhttt:t///// //gittt /iggtttt....... com zirui -ra-
Introduction
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 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 are-trained LMs leads to remarkable improvements in performance 77) As result, there is now an urgent necessity to effectively adapt these models, equipped with billion-scale parameters, to wide range of tasks.
However, 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
Equal contribution. The order of authors is determined by flipping coin.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

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-efficien tuning methods are proposed 8
13 [4]
Specifically, adapters [14 [3 insert 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 Although the optimizer states contribute to the memory footprint. storing activations (or
13 D2 feature maps) is the main memory bottleneck during training (often >> 70%) [16 19]. Thus,
parameter-effiiiet methods often do not reduce memory usage by much 9 19
In parallel, tuning. Since transformer-based models are mainly built based on the linear layer, a less-explored direction 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 A E R=xm and B E RTnx (in the context of machine learning. is often activations), we first sample k (k < m) column-row pairs according pre-defined distribution. Then we obtain A E Raxk and B' E R (k < m) by picking k columns of A and the corresponding rows of B according to the sampled column- now pairs [22]. Finally, we estimate AB R A'B'. In this way, we only need to store the sub-matrix A' and B' in GPU memory to perform the computation. Moreover, transformer-bon -aaed models training/tuning are performed with the first-order stochastic optimizer, e.g., Adam 115. 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.
reduce the main memory bottleneck by reducing the activation storage in fine-
The approximation method reduces the memory usage at the cost of giving outputs with variance. Thus there naturally exists an securacy-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 new family of unbiased estimator for matrix multiplication with reduced variance, dubbed Winner-Taie Take- All Column- Sampling (WTA-CRS) Compared to CRS, WTA-CRS reduces the variance of an estimator by focusing more high-probabilie regions of the sampling distribution. Morcover, WTA-CRS can serve as a drop-in replacement for the unbiased weight
80
Avenge Valibaiin
85- 80
(15--1--pp Large)
LattA (15-Laan))
LST (TS Langc)
Lelde-iT & CRCGR (TS Larget
LeRde-WT A-CSGtt. (TS
LARA- 1 (TS-p
LattA WTA- (19-3R) (R)
75
20
40
Training Peak Momory (GB)
Fig. 1. of WTA-CRS and other memory-efficien tuning methods. Unless specially stated, we use the T5-Large in the figure.
Accuracy- meerrryy memory trade-off
linear operation in transformers, providing gradient with reduced memory usage. As shown in Figure T method achieves better accuracy-memon trade-off than -tate-of-the-art memory-efficiae tuning methods, e.g., LST [9] and LoRA 12]. Moreover, since WTA-CRS executed the operation level, it orthogonal most of the existing parameter-effficient tuning methods. Our contributions are highlighted as follows:
We design 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 x peak memory reduction with almost accuracy drop, and enables up 6.4x 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 1100 GPUs with a capacity of 80GB. With WTA-CRS, it only requires 21. 21.6GB of memory for fine-tuning, which can run on GPU with 24GB memory, e.g. RTX3090Ti.
We implement WTA-CRS as ready-to-use extension for Pytorch with an easy-to-use API that can also be combined with other memory-saving techniques.
2

---

2 Background and Preliminary
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 mainly built based on the linear operation, which can be written as:
Menwey U'age
Medd Opeetin
86.1%
Forward Pass
z GEMM H, W);
(Ia)
88
Backward Pass
VH GEMM VZ, W
(1b)
(9-174)
04
VW- GEMM H
VZ),
(1c) Fig. 2. The GPU memory usage
breakdown for fine-tuning T5 [23]. where the batch size B is 64 and sequential length is 128 or 256.
where GEMM(.,)) -) is the General Matrix Multiplication operation, H and Z are the activation (or input feature maps) and output feature maps, respectively. W is the weight of the linear layer. VH, VW, and VZ are the gradient of H, W, and Z, respectively. From Equation (Tc), actívations H used in the backward phase. In commonly used deep learning framework [24] 25 it requires storing H GPU memory during the forward pass, for calculating the weight gradient VW in the backward pass.
Previous works show that although the model parameters contribute to the memory footprint, acti- vations (e.g., storing H) are the main memory bottleneck during training [16 I7 18 19. To get sense of the scale, we show in Figure 2 that for popular transformer models like T5, activations may take roughly 73 88% of the total memory, depending the batch size B and sequential length S.
2.2 Approximated GEMM With Sampling
Let X E RR Y E mx be two matrices. The goal is to efficiently estimate the matrix production XY. Singular Value Decomposition (SVD) outputs provably optimal low-rank estimation of X 21]. However, SVD is almost as expensive as matrix production itself. Instead, the sampling algorithm is proposed to approximate the matrix product XY by sampling & columns of X and corresponding rows of to form smaller matrices, which are then multiplied as usual 22 26):
GEMM(X, Y)= = Σ X Y
X kpi
i = X'Y',
(2)
where X E Rux1 and Y., E RRxx are the tth column and row of X and Y, respectively. In this paper, we call (X,,) , Y..)) the ,th column-row pair. l is the number of sampled pairs (1 < k S m). P
( pm a probability distribution over the column-row pairs. i, E {1.-.-m) m) is the index of the sampled column-row pair at the tth trial. is the scale factor. X' E B and Y R are the normalized sub-uutree matrices sliced according to the sampled column-row pairs.
Existing work 22 shows X'Y' is unbiased estimation of XY, i.e., E|X XY. Furthermore,
the approximation error [[[[X - X'Y'I|z] is minimized when the probabilities {p. { are proportional to the product of the column-row Euclidean norms 22 (Proof in Appendix
C
X l2 Y.
P; プ X. y
As we analyzed in Section 2.1 storing the activation H is the major memory bottleneck. If we can replace GEMM(H VZ) in Equation (Ic) with H' ZZ following the paradigm of Equation (2). then we only need H' instead of H in GPU memory to compute the gradient, which significantly decreases the memory usage of activations. This estimation linearly reduces the memory complexity from O(rm) to (nzk). Also, the total number of floating point operations (FLOPs) is reduced as
3

---

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-basee matrix multiplication.
3 Methodology
In recent years, have observed that deep neural network training can be performed almost entirely with first-order stochastic oplimization [15]. Thus intuitively, in stochastic optimization we can reduce the resources spent on oblaining gradients, long the estimated gradient is unbiased with reasonable variance [27] [28 229]. Following this motivation, we first design new unbiased estimator for matrix multiplication with reduced variance compared to the one in Equation (2) (Section (.16)
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 Taknerr-TakeAAll folwmn-Row Sampling). Following the notation in Section 2.2 let X E R Y E xm= be two matrices. P = {p:} 1T is the column-row index distribution in Equation We first define the variable f1) as f(1) = Y.
f(i) is an unbiased estimation for the matrix production between X and Y. To see this,
Eg~plf() Σ Pi Y XYY
m
We note that the prior approximated matrix multiplication in Equation ( the direct extension of f(i) by taking the average of {f(i) among k independent random trials to reduce the variance. Here we explore alternative approach reduce the variance of f(i) beyond simple averaging. Our core idea is to partition the column-row row index distribution P = } into two complementary regions based on the probability mass: a high-probabiling region pC and a low-probability region pPVc, where D {I...., 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;- We define WTA-CH CRS estimator for XY as follows:
E Bjuppic
Σ E((epe + (1 (1-Ep))(o)) LEC
(4)
GEC
We note that the random variable in Equation (4) is the column-row pair index j, and is only sampled from DVC. The estimator defined in Equation (4) contains two parts. The first part Evec fc)p, has no relationship with the random variable j and is summed deterministically. The second part f(j) is sampled stocastically, but scaled by the factor (1 - E Pc). When P {p.) ) concentrated on a small number of atoms, the scaling factor (1 Lee Pc) for the stochastic term should be small. Therefore, we intuitively expect the estimator have 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 unbiased estimator for matrix production XY, i.e E. B |E.Ec f())p + (1 - Pc) ff)l XXY
Theorem states that our proposed estimator in Equation is unbiased. Below we compare our proposed estimator to the CRS estimator in Equation ( 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 GEMM(X, Y) as:
&
(CRS)
g X.
Y
f(i),
6, ".....
Lid ~
P.
(5)
'Here we note that the theoretical analysis in this section can be applied to any probability distribution, not only limited the 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 C (IC) k), while stochastically average k [C] samples drawn from DCC to estimate the remaining terms, up to scale:
1 --Eecc
k-icl
(WTA-CRS) 6(X.)) = Et(rl(()+ +
Ltd Σ f(). 11, 1x-10 ~ PDic j=1
(6)
k-
oEC
Theorem 2 (Proof in Appendix
C.3
Suppose the total budget of column-row pairs k. IfC satisfies
C
Σ
P
P
k
AEC
then we have Varlg Y) Varlg(X, Y) Moreover, Varlg(X, Y)) is minimized when [C]
Lese
mee .k)
K-|c1
Both the left- and right-hand sides of Equation 7 depend on the size of the highest probability group [C], which controls the number of high probability column-row pairs that are directly added without sampling. Below we experimentally investigate whether Equation holds under the context of fine-tuning the transformer-based model with varying [C].
Query Layer 0.3
Key -.3()
Value Layer (k 0.3/2l)
d' Wi
0.5
0.8
Probablity Mass
0.66
layer
0.5-
0.4
6:1h layer
0.2
Raiio
0.2
0.2
0.7
Ratio
Rañi
Ratio
Fig. 3. The probability mass Exec Pe versus T in Equation (7) k: 0.3 I. Here we visualize
[CI
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.T
Experimental analysis. As shown in Figure 3 we visualize the two terms in Equation for the column-row index distribution of query, key, and value projection in the self-attention module, respectively [T. Specifically, we fix the total column-row pair budget k = 0.3/ and change the size the highest probability group |C| from 0 to k. We conclude that Equation 7 holds for most of the layers when munituning transformers. Thus, expect WTA CRS has better performance than CRS for adapting transformer- based models, which 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 de- scent 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 ap- proximated version in way that the estimated gradient is unbiased.
H
VW
VH
Equation (5)
Bubsample
W
Livel GPU
GRMM
Sabssmple
VZ
Z
Unbiasedness. Previous work has shown that ensure the unbiasedness of the gradient, the approximation can only be applied during the The rationale be- hind this conclusion is that we have [xf(x)] *
Fig. 4. The illustration of how to deploy WTA-CRS to linear layers. We only replace GEMM in Equa- tion Ich with its approximated version using WTA-CRS The pseudocode is given in Appendix D Algorithm
backward pass 16
[(Ear)) for any non-linear function /(-). e.g..
even when the approximation method gives an unbiased estimation, i.e., EH W)] HW Z, the output activations (e.g..
B[a²] & B²(ar), Thus if we replace the forward GEMM in Equation (a).
Gel.U(Z)) are still biased since the activation function is non-linear, namely.
GeLU( W)) (eZ)(E|Z) * E(GeeL())).
5

---

Next bleck Layers and
Drapout
Layer
Linear-O
Reshape
Dropont
Mul-2
Linear-D DD
1 [bxck K3x5]
GELU
1
SoftMax 1 [Poil K
Linear-V
Lincar-U
Ten-orMul-1
LayerNorm
K [Dx%
Linear-Q
Linear-K
Ouggrr [ex
block
Fig. 5. The diagram of single Transformer block. The shape of activations is annotated, where S.S. Dmddlt Nhead and Dread 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 acti- vation 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 unbiasness gradient and reduce the memory usage of storing H, as shown in the example of Figure 4 we only replace GEMM in the backward pass with its approximation (e.g., Equation (Ic)), while leaving the forward one unchanged (e.g., Equation (la)) We show in Appendix that the estimated weight gradient is unbiased in this case.
System Implementation. Here we present how we implement WTA-CRS in Equation (6) in practice. For the linear layer, as analyzed, only replace GEMM in Equation (Ic with its approximated version. In this case, X and Y in Equation (6) are activation H and output gradient VZ, respectively, Given the total column-row pair budget k the first step is to build the deterministic index set C, where each element is summed explicitly without sampling. Note that C is a set of indices with the highest probabilities in Equation (3). Thus, to build C, we only need to determine its size, denoted as C), which minimizes the variance of the estimator. As Theorem
2 suggested, we set
|C| mi, E-Eceg Pe The second step is to sample - C column-row indices from k-|C]
the remaining distribution PPIC obtain the set Canes where |Catol = k - IC). The third step is to build sub-sampled H' with only rows from C U Caree Note that for rows in H' from Cstoes need to
normalize by 1- L-Eese k k- according to Equation (6
We illustrate the above process in Figure 4
The pseudocode to Appendix D Algorithm
Scope. Here show which operation can be replaced with its approximation version. As shown in Figure 5) the transformer is mainly consisted of linear layer, TensorMul, and other operations (e.g.. GeLU, Dropout, LayerNorm). TensorMul in Figure(3refe four-dimensional tensors. Our WTA-CRA can be applied to Linear-Q. -K, -V. -0, -U, -D, TensorMul-1, and TensorMul-2 (in green). The activations of Dropout and GELU operations (in blue) can be losslessly compressed. The Softmax and LayerNorm operators (in gray) remain unchanged.
the multiplication between two
4 Related Work and Discussion
Due to the page limit, we discuss the related work on approximated matrix multiplication and aproximation in LLM inference. Other related topics, c.g., pfrfmettr-efficient fine-tuning, activation quantization, and gradient checkpointing, can be found in Appendix A and potential negative social impact in Appendix A
We also discuss the limitation
Approximated Matrix Multiplication. In the context neural networks, approximated matrix multiplccation methods can be broadly categorized into two main groups: (1) Butterfly-based methods replace dense weight matrices with butterfly matrices. We note that they focus on the weight matrix and are orthogonal to our research, as we concentrate on sub-sampling the activation matrix. (2) Column-row sampling (CRS) methods 22 31 select important rows and columns from the input matrices and perform the multiplication the sampled matrix. Our work is closely aligned
33 3

---

with this second research line. [1] share similarities with our research in terms of utilizing CRS for approximating matrix multiplication within neural networks. The main distinction lies in how to select the column-row pairs. Specifically, [2] deterministically selects column-row pairs without scaling, whereas our estimator divides the column-row pairs into deterministic component and stochastic component. As we analyzed, selecting column-row pairs deterministically is biased. Later we show that this approach may cause a significant accuracy drop ("Deterministic" in Figure 12).
Approximation in Inference There are two approximation techniques to reduce inference latency, (1) Sparse Modeling, which only involve subset of weights during computation reduce both computational and memory I/O demands B35
37 (2) Quantization, which compresses the trained 39
36
weight into lower numerical precision [8
40 All these approximation techniques trade-off model quality in return for improved efficiency. Besides approximation technique itself, researchers also propose methods to recover the accuracy drop of compressed models [4 42 e-g., by prompting compressed models. These methods can greatly improve the accuracy-effciency trade-off of LLMs.
Experiments
In this section, we design experiments to answer the following research questions: RQI: How effective is WTA-CRS in terms of accuracy with reduced memory usage? RQ2: How sensitive is WTA-CRS affected by its key hyper-parameters: RQ3: WTA-CRS contains two parts, i.e., the deterministic summation part and the statistical sampling part. Are they both necessary? RQ4: How is the fine-tuning speed affected by WTA-CRS?
5.1 Experiment Setup
Datasets and Evaluation Protocol. Following most of the previous work, adopt GLUE bench- mark 43 to evaluate the effectiveness of different methods, including the CoLA, SST-2, MRPC, QQP, MNLI, QNLI, RTE, and STS-B datasets. To evaluate the memory usage, we report the peak GPU memory usage and compression rate during the fine-tuning process with Huggingface API (44)
Compared Methods and Adopted Models. We consider three methods to compare in this paper: Full fine-tuning (Full), LoRA [12), and Ladder Side-tuning (LST) [9] Specifically, Full tunes all of the parameters in the model to provide upper bound of accuracy: LoRA inserts trainable low-rank matrices into the model to parameterize the weights' changes; LST injects a trainable ladder side structure. Since WTA-CRS essentially replace the linear operation with approximated one, emphasize that our WTA-CRS is compatible with all these three compared methods, i.e., they can be combined together towards smaller memory usage. For the backbone model, follow the previous work 9 T5-3B 446 for evaluating the effectiveness of different methods.
14 12
to adopt the OPT [45), Bert-Base (2), Bert-Large. T5-Base, LarLarge and
Hyperparameter Settings. For WTA-CRS, only has one hyperparameter k. which controls the column-row pair budget. We assign the same & to all replaceable linear operations in the model. We consider the normalized column-row pair budget k/ [D] E {0.3, 0.13. which are denoted as WTA-CRS @0.3 and WTA- WTA-CRS@O.1. respectively. We also consider the combination of WTA WTA-CRS and LoRA to further reduce the memory cost of both optimizer and activations. The detailed hyperparameters are given in Appendix F All reported results are averaged over three random trials.
5.2 Accuracy versus Memory Usage (RQ1)
To answer RQL, we first analyze the trade-off between the model performance and memory saving. The evaluation results and peak memory usage are given in Tables and respectively. We observe:
WTA- CRS achieves a superior trade-off between accuracy and memory usage compared baselines. Specifically, VTA-CRS has negligible accuracy drop, while the peak memory usage is reduced by 2.1x 2.7x (when combined with LoRA).
7

---

Table 1: The GLUE benchmark results with T5 and Bert at different scales.
Model
Mellod
CoLA
SST-2
MRFC
QQP
MNLI
ONLI
RTE
STS-B
AYG
00.91:1::
92.24.144
87.91.
87.81um
90.7
66.4
58.
BaRT-Dase
LORA
6166+a
91.7+a
900++03
86.9++1
83.6+tm0
90.8+0017
68 28
87.6+0
82.6
WTA-CRS@03
60.7
90,24b00s
87.0
87.5ta0
83 83.4
90.42207
65 65.92uu
89,3
81.8
LoRA-WTA-CRSE03
61,5221
89.62
89.6
86. 86.32
82
90.61a
67.9
87.3
Foll
60.12031
94.9222
91 91.32
88. Sthoo
87.02531
79:42.
90.622014
85.7
LoRA
60.6.am
94
92.2.
87 1..0s
86 2i.os
80.6
4
85.7
TS-Base
55.5+031
4+00
85.7+a15
93.4+00
72.7+nsi
90.4++00
WTA-CRS00.
60.9:20.52
94. 8dbo
91. 1:
88. Ottari
86. 3thom
93 1 1o07
78. 7:t..
90,5-boos
85.4
LoRA-WTA-CRSOO,
60.0:203:
94.42b0s
92.0
87.3 37.34
85.6
10 93.2
90.4
85.4
6688-011
72.6-m
90.2.40000
BERT-Large
LORA
65.9
93.8 17
90.8:b
87 6:booo
85 9ttoas
92.0maa 2
71 Sthoi
90. 3-.3mmmm
84:7
WTA-CRS@0.3
64.7:2044
93.5
89.5
88. 21:0
91.91:12
73.81
90.412u2
84.6
LRRA-WWA-CRS0033
66.0:a131
93.3
19.7
87 1::
91.94:
72.4-bmu7
89.7
Full
61.31211:
96. 96.311
93.4
89.7
89. 8.iu
94.2:
85 3
91.8.kuuu
87.7
LoRA
63.3
96.4401
93.5dmm
88 88.5
89 Stmas
94.3-oo17
84.2:bms
91.7
87.7
TS-Large
LST
59.9=031
95 8
91.8
88 48.4
7+t0s
94 94.2::002
82.5+an.8
91.4+bor
WTA-CRS@0.3
00.9-11
96.3++035
93.6++
89. 89.3+not
89.5+mn
94.1+
84.4+nsi
91.3++oa
87.4
LORA+WTA-CRSOUU
62.9:21:9
96 96.2
93.6r2047
88.32m02
89.2:tm0
94.0+2007
83.9
91.3-to00
87.4
T5.30
LORA
70.1+031
96.8+-13
94.0++3
89,9+50
0+00+
95 05.6+0005
85.9++%
92.9+om
LoRA-WTA-CRS003
71/4203
96/42b05
94.622m.9
90.0
91.0:2005
95.622012
86.34t
92.9
89.8
Table 2: Peak memory usage (GB) and compression rate of fine-tuning T5-Base and -Lrge. We measure the memory usage on single NVIDIA A100 (80GB) GPU. For T5-3B, since it is trained using multi-GPUs with data parallel. We instead report the maximum batch size in Figure for it.
UP
LaRA
LST
VTA-CRS@0.3
XTA-CRS@O.1
LoRA-KTA-CRSON
LoRA-WTA- CRSGO..
15-Base
17,66 (1x)
13.34((.3))
(3.2x)
(2
(2.4x)
6.50 (2.7x)
5.44 (3.2x)
T5-Large
45.85
(1.2X)
14.85
(2.1x
18.46 (2.5x)
17.44
14.64 (3.13x)
As we analyzed, LoRA mainly reduces the memory of optimizer states. Thus, although it has negligible accuracy drop, it can only achieve ~ 1.3x peak memory saving. LST can reduce memory us- age up to 3x, but its accuracy drop is much larger than LoRA and WTA-CRS. Since WTA-CRS executes at the operation level and fo- cuses the activations, to reduce memory usage more aggressively. When combing with LoRA, WTA-CRS achieves 2.7x memory usage saving with almost no accuracy drop. To fully tune T5-3B, it requires 66. 46B of mem- ory for in a mini-batch size of 32, which relies on A100 GPUs with a capacity of 80GB. On the other hand, LoRA+WTA-CON only requires 21.6GB of memory for finetuning with a mini-batch size of 32, which can run on GPU with 24GB memory, e.g. RTX3090Ti or A5000. We have experimentally confirmed this conclusion. Under the same hardware, WTA-CRS enables the tuning of larger models, resulting in improved down-streaming task performance. Thus as shown in Figure [ @ under the similar memory budget, WTA- -CRS outperforms the other methods in terms of the accuracy.
Park Men ond (GR)
8
Pxll
further combine LoRA with WTA-CRS
LaKtA
40
SRRX-WTA-RSBB
LARAI CRSB
100
Batch size
Fig. 6. Peak memory usage versus the maximum batch size of T5-3B. More similar results are shown in Appendix E.2
Also, according Figure 6 for T5-3B, LoRA itself can enable 1.9x larger batch size. When combined with LoRA, WTA - CRS enables 4.8x (k- - 0.. to 6.4x (k O.1 larger batch-size.
Influence of Row-column Pairs Budget (RQ2). As we analyzed in Section 3.2, WTA-CRS only have one hyperparameter, i.e., the total column-row pair budgets k. We conduct the ablation study with different budget k in Figure 7 We observe that 2 It has almost no accuracy drop when k = 0.3 DD) And the accuracy drop is about 1% when k = 0.1 (D). Notably, for T5-3B, the accuracy drop is only 0.4% when k- (.1 j, which is much smaller than T5-Base and T5-Large. This suggests that larger models are more compressible because they have more redundant activations, which is consistent with previous observations 47
5.3 Ablation Study (RQ3 and RQ4)
To answer RQ3, WTA-CRS is compared with two compositional methods to demonstrate its superiority. Namely, (1) the Deterministic method selects row-column pairs with top & probability of Equation (3)
We note that this is the estimator proposed in 211. (2) The CRS method follows Equation (3) to
8

---

Avtroge Valddtion Mealt e5
Atroge Valdmtie Kosih
85.4
954
Abrroge Valldiiin Reahh
87.9
57.7
67.4
89.5
85.0
67.1
36
58
82.5
Full
k|DD 0.3
80.0
86
LoRA
LoRA
w/LoRA
LoRA
(a) T5- Base
(b) T5-Large
(c) T5-3B
Fig. 7. Average validation results on GLUE dataset of STA-CRS with varying budgets.
sample the row-column pairs. All methods deployed GEMM in the backward pass, while leaving the forward one unchanged. The experiments are conducted on the training of T5-base language model on the SST2, MNLI, and QQP datasets; The column- row pair budget takes k/D) 0.1 for We observe:
all methods. The validation accuracy versus training epoch is given in Figure 12
@ NTA-CE CRS outperforms all compared methods, especially the training epoch grows. The deter- ministic selection of top k column-row pairs suffers from accumulation of bias error that ultimately results in failure of convergence. For CRS, also enables the unbiased weight gradient. However, as we theoretically and experimentally analyzed in Theorem 7 and Figure 3 it is worse than WTA-CRS due larger variance. In summary, both the deterministic and stochastic parts contribute to the effectiveness of WTA-CRS, which is consistent with our theoretical analysis.
The speed of WTA-CRS (RQ4). The configu- ration of computational infrastructure is given in Appendix F.I We note that WTA-CRS does not add any extra parameters to the model. Thus, WTA-CRS only affect the fine-tuning speed, without affecting the inference speed. Below we analyze how the fine-tuning speed affected by WTA-CRS- As we analyzed in Ap- pendix A limitation, the current implemen- tation is not heavily optimized and thus the execution time of WTA-CRS is still slower than the original linear operation (details are shown in Appendix E.2 However, under the same hardware, reduction in activation memory enables the of larger batch sizes, thereby improving training speed due to increased GPU utilization [48] 9 As we analyzed Figure WTA-CRS can enlarge the available batch size by up to 4.8x larger. This enhancement is expected to result in acceleration of the training speed. To illustrate, Figure [ presents a visualization of batch size against training throughput (sentences per second) for both T5-Large and T5-3B models. We observe that @ WTA-CRS enables faster training speed under the same hardware. Specifically, on the T5-Large model, WTA-CRS@O 0.1 shows 1.08x higher training throughput; and the T5-3B model, WTA-CRS @0.3 and WTA-CRS@E achieve 1.14x and 1.21 x higher training throughput, respectively.
14
Pull
Nu
50
150
(a) T5-Large
(b) T5-3B
Fig. 8. Batch size versus training throughput (sen- tences/sec) with different methods, where the quential length 128. The hardware is one single NVIDIA-A 100 (80GB).
6 Acknowledgements
The authors thank the anonymous reviewers for their helpful comments. The work is in part supported by NSF grants NSF IIS 184-885 IIS-2224843, and NSF Awards 2117439 and 2112606. This work made use of the High Performance Computing Resource in the Core Facility for Advanced Research Computing at Case Western Reserve University.
7 Conclusion
In this paper, we propose WTA-CRS, a new unbiased estimator for matrix production with reduced variance. We theoretically and experimentally show when and why the estimator is better than the traditional unbiased estimator in terms of the variance. In the context of adapting transformers, it almost has no accuracy drop while reducing the peak memory usage by up to 2.7x, and it enables 6.4% x larger batch size, which in return resulting in 1.2x higher training throughput.

---

References
[1] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in newral information processing systems, 30, 2017.
[2] Jacob Devlin, Ming-Woi Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04800 04800, 2018.
[3] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine I Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. The Journal of Machine Learning Research, 21(1):5485-5511, 2020.
[4] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1900 2020.
[5] Jingfeng Yang, Hongye Jin, Ruixiang Tang, Xiaotian Han, Qizhang Feng, Haoming Jiang, Bing Yin. and Xia Hu. Harnessing the power of Ilms in practice: A survey on chatgpt and beyond. arXiv preprint arXiv:2304.13712, 2023.
[6] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361, 2020.
[7] Ruixiang Tang, Yu-Neng Chuang, and Xia Hu. The science of detecting IIm-generated texts. arXiv preprini arXiv:2303.07205. 2023.
[8] Brian Lester, Rami AI-Rfou, and Noah Constant. The power of scale for parameter-efficient prompt tuning. arXiv preprint arXiv:2 (0.iv:2104008660 08691, 2021.
[9] Yi-Lin Sung, Jaemin Cho, and Mohit Bansal. Lst: Ladder side-tuning for parameter and memory efficient transfer learning. arXiv preprint (XXi::226600622.. 2022.
[10] Xiang Lisa Li and Percy Liang. Prefix-tuning: Optimizing continuous prompts for generation. arXiv preprint 1rXiv:2101.000 2021.
[11] Elad Ben Zaken, Shauli Ravfogel, and Yoav Goldberg. Bitfit: Simple parameter-efficient fine-tuning for transformer-bast masked language- models. arXiv preprint arXiv:2106./0199 2021.
[12] Edward Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. arXiv preprint arXiv:2106.00688. 2021.
[13] Rabech Karimi Mahabadi, James Henderson, and Sebastian Ruder. Compacter: Efficient low-rank hypercomplex adapter layers. Advances in Neural Information Processing Systems, 34:1022-1035. 2021.
[14] Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin De Laroussilhe, Andrea Gesmundo, Mona Attariyan, and Sylvain Gelly. Parameter-efficient transfer learning for nip. In International Conference Machine Learning, pages 2790-2799. PMLR, 2019.
[15] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 11Xiv:1412.6980, 2014.
[16] Jianfei Chen, Lianmin Zheng, Zhewei Yao, Dequan Wang, Ion Stoica, Michael Mahoney, and Joseph Gonzalez. Actnn: Reducing training memory footprint via 2-bit activation compressed training. In International Conference on Machine Learning, pages 1803-1813, PMLR, 2021.
[17] Paras Jain, Ajay Jain, Aniruddha Nrusimha, Amir Gholami, Pieter Abbeel, Joseph Gonzalez, Kurt Keutzer, and Ion Stoica. Checkmate: Breaking the memory wall with optimal tensor 2020.
rematerialization. Proceedings of Machine Learning and Systems, 2:497-511, 2020.
10

---

[18] Marisa Kirisame, Steven Lyubomirsky, Altan Haan, Jennifer Brennan, Mike He. Jared Roesch. Tianqi Chen, and Zachary Tatlock. Dynamic tensor rematerialization. arXiv preprint arXiv:2006.06666 2020.
[19] Muralidhar Andoorveedu, Zhanda Zhu, Bojian Zheng, and Gennady Pekhimenko. Tempo: Accelerating cansformer-basee model training through memory footprint reduction. arXiv preprint arXiv:2210. 10246, 2022.
[20] Yu-Neng Chuang, Guanchu Wang, Chia- Yuan Chang. Kwei-Herng Lai, Daochen Zha, Ruixiang Tang, Fan Yang, Alfredo Costilla Reyes, Kaixiong Zhou, Xiaoqian Jiang, et al. Discoverpath: A knowledge refinement and retrieval system for interdisciplinarity on biomedical research. arXiv preprint arXiv:2309.0 01000, 2023.
[21] Menachem Adelman, Kfir Levy, Ido Hakimi, and Mark Silberstein. Faster neural network training with approximate tensor operations. Advances in Neural Information Processing Systems, 31:27877-278888 2021.
[22] Petros Drineas, Ravi Kannan, and Michael w Mahoney. Fast monte carlo algorithms for matrices i: Approximating matrix multiplication. SIAM Journal on Computing, 36(1):132-157, 2006.
[23] Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. Glue: multi-task benchmark and analysis platform for natural language understanding. arXiv preprint trXiv:1804.0746/1 2018.
[24] Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, al. Tensorflow: system for large-scale machine learning. In Osdi, volume 16, pages 265-283. Savannah, GA, USA, 2016.
[25] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
[26] Petros Drineas and Ravi Kannan. Fast monte-carlo algorithms for approximate matrix multi- plication. In Proceedings 42nd IEEE Symposium on Foundations of Computer Science, pages 452-459. IEEE, 2001.
[27] Brian Chmiel, Itay Hubara, Ron Banner, and Daniel Soudry. Minimum variance unbiased M sparsity for the neural gradients. In The Eleventh International Conference on Learning Representations, 2023.
[28] Brian Chmiel, Ron Banner, Elad Hoffer, Hilla Ben Yaacov, and Daniel Soudry. Logarithmic unbiased quantization: Simple 4-bit training in deep learning, arXiv preprint arXiv:2112.10769, 2021.
[29] Deniz Oktay, Nick McGreivy, Joshua Aduol, Alex Beatson, and Ryan P Adams. Randomized automatic differentiation. arXiv preprint arXiv:2007.10412 2020.
[30] Xiaoxuan Liu, Lianmin Zheng, Dequan Wang, Yukuo Cen, Weize Chen, Xu Han, Jianfei Chen, Zhiyuan Liu, Jie Tang, Joey Gonzalez, et al. Gact: Activation compressed training for generic network architectures. In International Conference on Machine Learning, pages 14139-14152. PMLR, 2022.
[31] Zirui Liu, Shengyuan Chen, Kaixiong Zhou, Daochen Zha, Xiao Huang, and Xia Hu. Rsc: Accelerating graph neural networks training via randomized sparse computations. arXiv preprini arXiv:2210.10737, 2022.
[32] Muralidhar Andoorveedu, Zhanda Zhu, Bojian Zheng, and Gennady Pekhimenko. Tempo: Accelerating transformer- based model training through memory footprint reduction. arXiv preprint arXiv: 10246, 2022.
11

---

[33] Beidi Chen, Tri Dao. Kaizhao Liang, Jiaming Yang, Zhao Song, Atri Rudra, and Christopher Re. Pixelated butterfly: Simple and efficient sparse training for neural network models. arXiv preprint arXiv:2/ /2.00029, 2021.
[34] Tri Dao, Beidi Chen, Nimit S Sohoni, Arjun Desai, Michael Poli, Jessica Grogan, Alexander Liu, Aniruddh Rao, Atri Rudra, and Christopher Ré. Monarch: Expressive structured matrices for efficient and accurate training. In International Conference Machine Learning, pages 4690-4721. PMLR, 2022.
[35] Shaochen Zhong, Guanqun Zhang, Ningjia Huang, and Shuai Xu. Revisit kernel pruning with lottery regulated grouped convolutions. In International Conference on Learning Representa- lions, 2022.
[36] Elias Frantar and Dan Alistarh. Sparsegpt: Massive language models can be accurately pruned in one-shot. 2023.
[37] Zichang Liu, Jue Wang, Tri Dao, Tianyi Zhou, Binhang Yuan, Zhao Song, Anshumali Shrivas- tava, Ce Zhang, Yuandong Tian, Christopher Re, et al. Deja vu: Contextual sparsity for efficient Ilms at inference time. In International Conference on Machine Learning, pages 22137-22176. PMLR, 2023.
[38] Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. Llm. int8 O): 8-bit matrix multiplication for transformers at scale. arXiv preprint arXiv:2208.0777 2223222388.73,,,... 2022.
[39] Ji Lin, Jiaming Tang, Haotian Tang, Shang Yang, Xingyu Dang, and Song Han. Awq: Activation-aware weight quantization for Ilm compression and acceleration. arXiv preprins arXiv:2306.00998, 2023.
[40] Guangxuan Xiao, Ji Lin, Mickael Seznec, Julien Demouth, and Song Han. Smoothquant: Accurate and efficient post-training quantization for large language models. arXiv preprint arXiv:2211.10438 2022.
[41] Zhaozhuo Xu, Zirui Liu, Beidi Chen, Yuxin Tang, Jue Wang, Kaixiong Zhou, Xia Hu, and Anshumali Shrivastava. Compress, then prompt: Improving accuracy-efficiency trade-off of Ilm inference with transferable prompt. arXiv preprint wrxiv:230.111166 2023.
[42] Mingjie Sun, Zhuang Liu, Anna Bair, and J Zico Kolter. A simple and effective pruning approach for large language models. arXiv preprint arXiv:2306.
1695, /695,, 2023.
[43] Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding. arXiv preprint arXiv: 1804.07461, 2018
[44] Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, al. Transformers: State- of- natural language processing. In Proceedings of the 2020 conference on empirical methods in natural language processing: system demonstrations, pages 38 23-45,,2 2020.
[45] Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. Opt: Open pre-trained transformer language models. arXiv preprint arXiv:2205.0100 2022.
[46] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with unified text-to-text transformer. The Journal of Machine Learning Research, 21(1):5485-555 5551,, 2020.
[47] Zhuohan Li, Eric Wallace, Sheng Shen, Kevin Lin, Kurt Keutzer, Dan Klein, and Joey Gon- zalez. Train big, then compress: Rethinking model size for efficient training and inference of transformers. In International Conference on machine learning, pages 5958-5968. PMLR, 2020.
[48] Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in I hour. arXiv preprint 1rXiv:1706.02667, 2017.
12

---

[49] Beidi Chen, Zichang Liu, Binghui Peng. Zhaozhuo Xu, Jonathan Lingjie Li, Tri Dao, Zhao Song, Anshumali Shrivastava, and Christopher Re. MONGOOSE: learnable LSH framework for efficient neural network training. In International Conference on Learning Representations (ICI.R), 2021.
[50] Zirui Liu, Kaixiong Zhou, Fan Yang, Li Li, Rui Chen, and Xia Hu. EXACT: Scalable graph neural networks training via extreme activation compression. In International Conference Learning Representations, 2022.
[51] Guanchu Wang, Zirui Liu, Zhimeng Jiang, Ninghao Liu, Na Zou, and Xia Hu. Towards memory efficient training via dual activation precision. arXiv preprint arXiv:2208.04110 2022.
[52] Guanchu Wang, Zaid Pervaiz Bhat, Zhimeng Jiang, Yi-Wei Chen, Daochen Zha, Alfredo Costilla Reyes, Afshin Niktash, Gorkem Ulkar, Erman Okman, Xuanting Cai, et al. Bed: A real-time object detection system for edge devices. In Proceedings of the 31st ACM International Conference on Information & Knowledge Management, pages 4994 4998, 2022.
13

---

Appendix
A Extended Related Work and Discussion
Sarameter-Efficieen Fine-tuning. Parameter-cfficient tuning methods select small subset of parameters or insert few parameters pre-trained network. Then they only update the small subset of parameters, while keeping others fixed [
14 49] For example, Adapters insert a small module into the transformer blocks and only update it. Similarly, prompt tuning g introduces a small vector that is concatenated with the imput embeddings. BitFit [1]] only tunes the bias term of the model. LoRA [12] injects trainable rank decomposition matrices into the transformer block. Although these methods are "parameter-efficiini
9 T0 I
14
they actually cannot reduce the memory usage of the model itself. This is because we still needs build the computation graph for the whole model. Instead, the memory usage of optimizer states will be significantly reduced, which in proportional to the number of trainable parameters [5]
Activation Quantization. The activation quantization methods focus on quantizing the activation
into low numerical precision numbers, e.g., 8-bit integers [16, 30 50 [51 52 Our work also compresses the activation, but in a different way. We emphasize that our work is orthogonal to activation quantization in the sense that our work essentially reduces the dimension of activation. This distinction allows our method to be readily combined with activation quantization techniques, offering the potential for even more aggressive compression.
Gradient Checkpointing. Gradient checkpointing helps decrease activation memory usage by saving only a selection of activations. However, it demands additional computation during the backward pass, as discarded activations must be recalculated [18. [7]. According to the report of Checkmate [T)), it achieves "a 2.3x memory reduction when training a BERT model with Checkmate optimizations (at Ix extra overhead for rematerialization)..
Limitations Although WTA-CRS significantly reduces the computation of the backward pass in a hardware-friendly way i.e., dropping entire rows/columns in the tensor, the current implementation still hampers the execution time of linear operations. This is because the extra sampling process and data movement counteract the acceleration. However, we note that (1) the overhead can be greatly reduced with better implementation, e.g., using prefetch and operation-fusian technique [30; (2) the existing implementation can still yield large speedup when employing larger batch sizes (Figure R).
Potential Negative Societal Impacts. Our research primarily focuses on reducing the memory requirement of fine-tuning Language Models (LMs). The carbon emissions produced by LM fine- tuning may pose environmental issues. Our next step is to further improve the efficiency of LM fine-tuning, particularly on hardware with lower energy consumption.
B Unbiasedness of Weight Gradient
This part we directly follow the proof of Theorem 1 in ActNN [6]. For completeness, we provide the proof sketch here that is short and easy to follow. Specifically, here we use ReL.U the activation function for illustration convenience. We note that the conclusion in this section holds for any non-linear activation function. Specifically, the forward pass of ReLU-Linear at the 1 layer
11+ H()())
H (1+1) = (e1)
and the backward pass of ReLU is:
E (+1)
E z > VH (1)
I | ① (1 [+]]
14

---

where @ is the element-wise product and is the indicator function. The element-wise product is linear operation and A is only related to the pre-activation ( in the forward pass. We only apply the approximation during the backward pass so Iza+1>s0 can be extracted from the expectation. We know that for the last layer L, we have ([V H since we do not
apply activation at the output layer. We then can prove by induction that E| H (!-1) and w for any layer I.
C
Proof
C.1 Derivation of Equation
Let X E ERRM Y E Rmxx be two matrices. The matrix multiplication XY can be estimated as
GEMM X. Y)
X be Y.. = X'Y'
kpi
t=1
Equation
3 shows the approximation error E
X
Y |el] is minimized when the probabilities
X la
Y
Yiell
Pr
1=
Y 3
Proof. Let [(i) Y.
E RRxy We note that f(i) is an unbiased estimation of XY, Namely,
P:
m
pi X
Bx-D[(()
Σ
XY.
Then have
h X'Y' f(i),
(8)
where i1,". ......, is are the index of the sampled column-row pairs at tt random trials. For each é, its variance is
X
Y
Var(f(t,)
Var
Pi
X
Y 1v
E
E
P
Pi.
Y?
E
(XY)² 2
P
X2
Y
(XY)
Pt
E=1
where the first step follows from the fact that Var
B -
Then we have,
E X XY
XY Σ E X
j=l
Var (ar(('
15

---

By combining Equation g and Equation 9
into the above equation, we have
XY
X'Y [1]
X
Pr
XX F.
Pt
To minimize X - X r]. the optimal probability distribution can be obtained via solving the following optimization problem:
Y
mii
Pt
m
s.t.
Pt
The solution the above convex problem is the distribution defined in Equation
Namely,
+ l2
Y.. |Y
Pi
2==1
m X la
gall2
C.2 Unbiasedness of Our Proposed Estimator
Theorem 1 (Proof in Appendix C.2 The estimator defined in Equation (4) is an unbiased estimator for matrix production XY, Le, E. Eaueore 1EeEc (c)pe + "ecc Pe [ XY.
Proof.
E Ejppoie Σ f(c)pe + (1 LL1--pol +11--x))6i) Epdf(()
E(eepe + +1 (1- Σ
Pe) (xppeepe)
GEC
OEC
Σ f(c)pp +(1 15--pe Pc) Σ
Pj
f(j)
1 --Eecc Po
SEC
oEC
JEDYC
Σ f(c)pe+ Σ f(iis;
CEC
JEDYC
=ByDDl((i
-Y
C.3 Variance of Our Proposed Estimator
Theorem 2 (Proof in Appendix C.3
Suppose the total budget of column-row pairs is k. IfC satisfies
C
Σ
Pe
k
CEC
then we have Varlg(X, Y)] Varig(X, Y)). Moreover Varlg(X, Y)] minimized when [C]
E.ce
min .2) [C
16

---

Proof. Recall that the original estimator for matrix production X is defined
Bp)lf(()
(10)
and our proposed family of estimator is defined as:
h(j) = E juppie
Σ
f(e)pe + (1- -PDDfff()
(11)
AEC
aEC
We first define three independent random variables as belows:
u pe,
(12)
j ~ PPIC
(13)
b ~ Bernoulli(1
Σ Pc).
(14)
OEC
According to the Law of total variance, we have
Var[f(i)] Es
Varlf(0) b] + Var [[
E. Var [()]6]]
Σ Pe Var(f( 0) +1 Σ
pe) De Var [ar[s(t)]b 1]
CEC
2
Σ (a) DIC] oEC
(15)
where the first step follows from the fact that for any random variance x, y, we have Varly]
[[Var[y[x]] + Var[E[y|x]]. Also, by Equation 11). we have
Var[4(]] = (1 Σ
p E()
(16)
CEC
By combining the above two inequality, we have
Var[h(G)] s (1
Σ pe) Varl
(17)
CEC
Equation 17 quantitatively shows the variance reduction of h(j) over f(i). Then compare estimator g [ Y) and g(X, Y) in terms variance.
First, because g(X, Et f(tt). i1.-..ig
in lid
P. Thus have
Varlg(X, Y)! T Var[s(ii]-
(18)
Similarly, we have
Varlg Y)] k-
Var[h(j)]
(19)
By combining Equation
into the above equations, we have
Varlg(X, (X, Y)
Var[h(i)]
(20)
k C
< Enec Po Var[f(i)]
CC)
Ecec Pe
/arlll Y)).
C
17

---

where the first step follows from Equation (17 By setting 1-E k arrive the conclusion k-1c
that when Ecec Pe we have Var [ X Y) s Varlg(X, Y)).
L-Zege Pe k achieves the minimal when min clefo .
1- Eeec Pe.
Further,
k- [2]
k- C-
D Implementation Details
The pseudocode for approximated linear layer with WTA-CRS and standard line layer is given in Algorithm and Algorithm respectively. The column-row row pair sampling procedure is given in Algorithm For the ease of illustration, we ignore the sequential length. As we mentioned in the main text, we only replace the GEMM the backward pass with WTA-CRS According to Equation (Tc). we need the activation gradient VZ to perform the column-row pair sampling during the forward pass. Thus we initialize cache in CPU memory to store the gradient norm of activations from the last step. When performing column-row pair selection, between CPU and GPU, which will cause extra time overhead due to the data movement. Fortunately, we note that the number of elements in the gradient norm of activations is significantly less than the one in activations, which does not cause significant time overhead.
need to swap the gradient norm of activations
Algorithm 1: Forward & Backward pass of Approximated Linear Layer Hyperparameter: The total budget of column-row pairs k. procedure INIT:
Initialize Cache E R as an empty matrix in main memory // N is the total number of samples the dataset. Cache output gradient VZ.
used for saving norm of
end
procedure FORWARD PASS:
Input: activation H BxD, weight W E PxD, indices of the current batch samples BI = [i,,... ,jn].
ctx <- 0 // the context which saves tensors for backward
= HW
H', ind SUUBAALPLLHH Cache[B[]] k)
// Cache [BJ] is the cached gradient norm from the backward pass; ind is the set of involved column-row pair indices ctx < (H', BI, return Z
end
procedure BACKWARD PASS:
Input: ctx from the forward pass, output gradient VZ E B
H'. W. BI, ind f- ctx
VH = VZW
VZ < VZ(ind
// Z' E &
VW H'VZZ'
for in BI do Cache[j]= 1 end
// Update the gradient of samples in the current batch return VH, VW
end
18

---

Algorithm 2: SUBSAMPLE
Input: activation H E R
BxD,
gradient norm E IR the total budget of column-row pairs k.
for 1, 1,... B
do
2. H..lla
// The probability of column-row pairs defined in
Σj= 2. H
Equation 3
end
Le s.t. C = [k]. // C is the set of column-row pair (e.... k-
indices associated with [C] largest P/-
Sample k - i..d column-row pairs Cstoe = (ir, d from the distribution PDlc ind Cste
for j E Croe do
j H[j.:]
1-E [k Chip
// We need normalize the stochastic part (6) to ensure the unbiasedness.
in Equation
end
H' < H(ind] return H', ind
// H' E R
Algorithm 3: Forward & Backward pass of the standard Linear layer
procedure FORWARD PASS:
Input: activation HQ E IR BSXD weight Wq E DxD, batch indices inder ctx {} // the context which saves tensors for backward
Zo Ho W Q
ctx H Wo}
return Zo
end
procedure BACKWARD PASS:
Input: ctx from the forward pass, output gradient VZQ
Ho, W + ctx
VHO VZG W
VW Q VZQ
return VHQ, VWQ
end
19

---

E More Experimental Results
Query Layer (k=0, (2)
Key Layer (& D.)
Value Layer (2 (2-00.] 0.1 PD)
00.6
byyer
0.6
0.5
4-th boyer
Prolabiity Maxe
6-th byyer
0.4
0.4
Ratio
0.1
0.2
0.3
0.4
0.5
0.6
0.2
0.5
0.3
0.4
10
Rario
Ratio
1d
Ratio
Equation at = 0.1 Here we visualize the column-row index distribution of query/key/value layer T5 base model, fine-tuned on RTE dataset.
Fig. 9. The probability mass Eesc Pe versus I
Query Layer (k p.5
Key Layer 0.5)D)
Value Layer (K D.)
80.75
0.75
0.75
Probebiity Maxe
2
0.50
4-th layer
0.50
0.50
0.2
Ratio
0.25
0.25
0.1
0.2
0..
044
0.5
0.6
07
011
0.3
0.4
0.6
0.7
011
0.2
0.3
0.4
0.5
0.5
0.7
Id
Ratio
Rtttio
Fig. 10. The probability mass Eeeg Pe versus in Equation 7
[C]
at l: 0.5 D) Here visualize
the column-row index distribution of query/key/value layer T5-base model, fine-tuned on RTE dataset.
Query Layer (22 0.1)
Key Layer
0.1)
Value Layer ( 0.1)
2-th byyer bayer
0.7
Probabiity Mao
6
0.6
0.6
hwn
Mu
hum Map
Maw
Nieration
Fig. 11. The probability mass of top-10% column-row roww pairs in Equation versus iterations. Here
we visualize the query/key/value layer T5 bS-b model, fine-tuned on RTE dataset.
Accuney (%)
94
Accunscy (%)
80
(6)
80
Desemirsatic
92
70
60
WTA-ORS
10
2
10
2
$
10
(a) SST2
(b) MNLI
(c) QQP
Fig. 12. Validation results of T5-Base with different methods.
E.1 More Experimental Analysis on Theorem
To evaluate Theorem more comprehensively, below we also plot the Σ (.1 and = 0.5 in Figure 9 and 10 respectively. We also plot Eeee Pe versus iterations in Figure 1T We summarize that The the column-row index distribution is concentrated a few column-row pairs. Thus, the assumption Theorem holds under the context of fine-tuning transformers.
Lesc versus [C] F at k
E.2 More Experimental Speed Analysis
Increasing the batch size can often result in faster model convergence and/or enhance the final performance. Ideally, should adjust the batch size according to the requirements of our model
20

---

T5-
T5-
T5
T5-
Method
ATT
FF
Block
Large
Full
8
10
17
1052
Fwd
WTA-CRS
22
16
37
2013
Full
16
19
34
2073
Bwd
WTA-CRS
15
14
30
1738
Full
24
29
51
3125
F-B
WTA CRS
37
30
67
3751
Table 3: Latency (ms) of Forward and Backward pass.
rather than being constrained by the GPU's memory capacity. To illustrate this, we have represented the correlation between peak memory usage and maximum mini-batch size for T5-Base T5-Large, and T5-3B in Figure I3 Our observations highlight that WTA-CRS effectively increases the maximum available batch size.
We also provide the apple-to-apple speed comparison for linear operation with and without WTA-CRS in Table 3 In Table "Fwd", "Bwd", and "F-B" are the time of forward pass, the time of backward pass, and the total time for both the forward and backward pass, respectively. We summarize that under the same workload, the current implementation of WTA-CRS may roughly slow down the linear operation about 20%. This is because the extra sampling process and data movement counteract the acceleration (see Algorithm [ However, we note that (1) the overhead can be greatly reduced with better implementation, e.g., using prefetch and operation- Peeaation-fusion technique [30); (2) the existing implementation still yield large speedup when employing larger batch sizes (Figure X
80
30
801
Prok Mem cont (CBB)
Pauk (CBB)
Prok Moen cost (GBB)
60
600
40
40
40
20
LaRA-WTA-CRE WTA-CRSGO
20-
500
1000
1500
200
400
500
100
200
Butch size
Batch size
Butch size
T5-Base
(b) T5-L.arge
TS-3B
Fig. 13. Peak memory usage versus maximum mini-bater size of T5.
E.3
Performance of WTA-CRS on OPT models
We conduct experiments of WTA-CRS OPT models. It is observed in Table 4 that WTA-CRS cxhibits almost no drop in accuracy compared to Full training and LoRA on all three architectures. This observation indicates the effectiveness of WTA-CRS, especially considering its consistent performance across the encoder-only, decoder-only and encoder-decoder transformer architectures.
Method
CoL.A
MRPC
RTE
STS-B
Average
Full
49.84
85.47
72.56
84.43
73.075
LoRA
52.3
88.36
74.01
87.21
75.47
LoRA+WTA-CRS@033
51.8
88.43
74.01
86.61
75.2125
Table 4: Deployment of WTA-CRS@O.3 to OPT-350M.
E.4 Performance of Activation Quantization+iiT--CR
To demonstrate the orthogonality of WTA-CRS with activation quantization, we conducted additional experiments of combining WTA-CRS @0.3, along with activation quantization @ 8bit, the T5-
21

---

base model. The experiment results are presented in Table 5 It is notabe that the combination of MORR+WTA, SEFTA-CRSOO PRRS@O.3+uuam @bbi exhibits almost no drop in accuracy compared with Full training, LoRA, WTA-CRS@033 and the orthogonality of WTA-CRS with activation quantization, indicating that they can effectively applied together without compromising performance.
This observation clearly demonstrates
Method
CoLA
MRPC
RTE
STS-B
Average
FP32
60.1
91.5
79.4
90.6
80.4
LoRA
60.6
92.2
80.6
90.7
81.0
WTA-CRS@0.3
60.9
91.1
78.7
90.5
80.3
LoRA+WTA-CRS@0.3
60
92
80.1
90.4
80.6
LORA-WTA-CRSS003-000 Sbit
60.3
92.06
81.2
90.4
81.0
Table 5: Combination of WTA-CRS with quanzation on the T5-base model.
E,5
Comparison WTA-CRS with QLoRA
We conduct experiments to compare WTA-CRS with QLoRA on the T5-Base model. Specifically, we first quantize the model into NF4 data format using 'bitandbytes' and keep it fixed. Then we apply LoR LORA+WTA-CRSER 0.3 over the quantized model. Below in Table we demonstrate that applying QLoAA and -TT- over base model quantized in NF4 results in almost no performance loss against the QLoRA baseline. However, such joint applications may enjoy the training-time memory efficiency offered by WTA-CRS, as illustrated in Table 7
Method
Cola
MRPC
RTE
STS-B
Average
LoRA
60.6
92.2
80.6
90.7
81.025
LORA-WTACRS@0.3
60.0
92.0
80.1
90.4
80.625
QLoRA (NF4)
61.3
92.2
81.2
90.5
81.3
QLoRA WTACRSSO
62.1
92.2
82.7
90.1
81,775
Table 6: Comparison of WTA-CRS with QLoRA.
Method
Base
Large
Lora
13.84
36.83
QLora
13.64
36.
LORA-WTA-CRS00.3
6.50
17.44
QLORA-WTA-CRSS0.3
6.31
16.75
Table 7: Peak memory usage (GB) of fine-tuning T5-Base with different methods.
Experimental Settings
We give the detailed hyper-paraae arprr-paraeeter setting in this section. Specifically, for both T5 and BERT models, the parameters are updated with the Adam optimizer 31 = 0.9 3g 0.999 e 10- and weight decay 0. The the learning rate is adjusted with a linear LR Scheduler, which maintains constant learning rate for the initial 500 steps, and adjusts it gradually thereafter. The input sequences are padded to the maximum length 128. WTA-CRS has LoRA dimension 32 it is combined with LoRA. To achieve the optimal solution, the T5-Base, Large, 3B and BERT-Base and Large models have different learning rate, training epoch number, and mini-batch size on different datasets, which respectively.
given in Tables 9 10 11
F1 Computational Infrastructure
The computational infrastructure information given Table g
22

---

Table 8: Computing infrastructure for the experiments.
Device Attribute
Value
Computing infrastructure
GPU
GPU model
NVIDIA-AI00
GPU Memory
81251MB
CUDA Version
11.4
CPU Memory
512GB
Table 9: Learning rate.
Model
Method
CoLA
SST-2
MRPC
QQP
MNLI
QNLI
RTE
STS-B
WTA-CRS@0.3
2c-5
BERT-Base
LoRA+WTA-CRSGN
2e-4
5e-4
2e-4
3e.4
3e.4
2e.4
3e-4
WTA-CRS@0.3
3e-5
3e-6
3e-5
3e-5
WTA-CRS@O.1
3e-5
T5-Base
LoRA-WTA-CRS@0.3
3e-4
3e-5
3e-4
3e-5
3e-5
3e-5
3e-4
3e-4
LoRA+WTA-CRS@0.1
3e-4
3e-5
3e-4
3e-5
30-5
3e-5
3e-4
3e-4
WTA-CRS@0.3
2e-5
BERT-
LoRA+NTA-CRS@003
3e-4
2e-4
2e-4
2e.4
2e.4
2e.4
3e-4
3e-4
VTA-CRS@0.3
3e-5
3e-6
3e-5
3e-5
TS-Large
WTA-CRS@0.1
3e-5
3e-6
3e-5
3e-5
LoRA+WTA-CRSGH
3e-4
3e-5
3e-4
3e-5
3e-5
3e-5
3e-4
LoRA+WTA-CRS@0)
3e-4
3e-5
3e-4
3e-5
3e-5
3e-5
3e-4
LoRA+WTA-CRS@O3
3e-4
3e-5
3e-4
3e-4
3e-4
3e-5
3e-4
3e-4
15-38
LoRA+NTA-CRS60011
3e-4
3e-5
3e-4
3e-4
3e-4
3e-5
3e-4
3e-4
Table 10: Training epoch number.
Model
Method
CoL.A
SST-2
MRPC
QQP
MNLI
QNLI
RTH
STS-B
WTA-CRS@0.3
20
20
10
10
10
10
20
10
BERT BERT-Base
LoRA-WTA-CRS@0.3
60
20
20
20
20
20
40
40
WTA- TTA-CRS00.3
40
10
20
10
10
10
50
20
WTA-CRS@0.1
40
10
20
10
10
10
50
20
T5-Base
LoRA+*TA-CRS@003
40
20
20
20
20
LoRA+*TAACCB660.1
40
20
20
20
20
WTA-CRS@0.3
60
20
20
10
10
10
40
10
BERg-Large
LoRA+WTA-CRS@0.3
60
20
20
20
20
20
40
40
WTA-CRS@0.3
20
20
10
10
10
VTA-CRS@001
20
20
10
10
10
20
Trrga
1.
40
40
10
10
10
20
LoRA-WTA-CRS@O.I
40
10
20
10
10
10
60
20
LoRA-WTA-CRS@003
40
20
10
10
10
60
20
T5-3B
LoRA+WTA-CRS@O.1
40
20
10
10
20
Table 11: Training mini-batch size. Here for BERT models, utilize 128 batch size except STS-B, where a batci size 16 is used. For T5 model, we use 100 batch size across all datasets.
Model
Method
CoLA SST-2
MRPC QQP MNLI ONLI RTE STS-B
WTA-CRS@O0
128
16
BEsl-Base/Large
LORA4TTA-CRS@0
128
16
100
WTA-CRS@0.
100
T5- 35-Base/T.arge//3
-ORA+VTA-CRSB
100
LORAATTA-CRS@0
100
23


Reviews:
[{"review_id": "LBdtus6Tic", "reviewer": "Reviewer_ky3t", "summary": "The paper's contribution is in proposing a practical, intuitive yet not trivial unbiased approximation to gradient training of matrix multiplication. It shows that even though totally deterministic sampling is biased, somewhat deterministic sampling is unbiased, and a judicious allocation of sampling to those pairs favored by deterministic thinking can lead to the use of a larger batch size with empirically negligible performance loss. This reviewer must declare that he does not check the derivation very carefully.", "questions": "This reviewer does not have a question.", "limitations": "N/A", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The proposed idea is practical and can be readily combined with virtually all first-order gradient-based training methods.\nThe paper also derived why deterministic sampling is a biased estimator and empirically shown the associated bad performance, thus proving that the additional complexity of stochastic sampling over deterministic sampling is not only sufficiently better but also necessary.", "weaknesses": "It's just a few empirical comparisons, but the performance gap between CRS and WTA-CRS seems modest.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for spending time and effort in reviewing our paper. We appreciate your constructive comments and suggestions for improving the quality of this work. The feedback truly encouraged us to increase our efforts in conducting quality and impactful research."}, {"review_id": "GuqeTX6P9u", "reviewer": "Reviewer_cMiu", "summary": "The authors studied fine-tuning LLMs with limited memory. As the increased scale of current LLMs, the memory cost during fine-tuning is of great importance when adapting the pretrained LLMs to down-streaming tasks. In contrast to the existing work that mainly focus on the number of updated weights, this paper proposed to reduce the number of stored activations, also the inputs to each layer. Given the widely used stochastic gradient descent optimization pipeline, the authors proposed to store a subset of activations that can generate an unbiased gradient estimation. This way, the training memory and the training time decreased significantly. The authors provide both theoretical and experimental analysis on their CRS methods. ", "questions": "Overall, I think this paper has a relatively high quality in both writing and scientific contribution.", "limitations": "Yes", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- This paper studied an important problem in LLM fine-tuning, i.e., how to fine-tuning LLMs with less memory consumption without increasing the computation cost. The authors provided solid quantitative results to show that the main memory consumption is from storing the intermediate activations. \n- The authors provided a general solution for fine-tuning LLMs under memory constraints. The solution can be applied in most transformer-based network architectures.  \n- The authors provided solid mathematical proof on the unbiased gradient estimation, which is especially encouraged. \n- The extensive experiments on different network architectures showed the efficacy of the methods.\n- The released code can benefit the following researchers studying efficient LLM fine-tuning. ", "weaknesses": "- I am not fully convinced by the comment made in Line241-244, i.e., the methods in the paper is orthogonal to the activation quantization. When activation is quantized into a lower bit width, it is very possible that the number of less important activations will decrease. This way, the selection on the top-k columns in activation matrices with the proposed methods may hurt the training accuracy or convergence. It would be great if the authors can provide some theoretical analysis or experimental results on this combination. Otherwise, it would be necessary to provide some comparison results w.r.t. the activation quantization.\n- It would be great if the authors can discuss the main difference of their paper w.r.t. [Randomized Automatic Differentiation, ICLR2021].\t ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**[W1] Whether WTA-CRS is compatiable with activation quantization or not** (experiment done)\n\n\nWe thank the reviewer for this thoughtful comment. We acknowledge the importance of demonstrating the orthogonality of WTA-CRS with activation quantization. To address this concern, we conducted additional experiments of combining WTA-CRS@0.3, along with activation quantization@8bit, on the T5-base model.\n\nThe experiment results are presented in the following table. It is notabe that the combination of LoRA+WTA-CRS@0.3+Quant@8bit exhibits almost no drop in accuracy compared with Full training, LoRA, WTA-CRS@0.3, and LoRA+WTA-CRS@0.3. This observation clearly demonstrates the orthogonality of WTA-CRS with activation quantization, indicating that they can be effectively applied together without compromising performance.\n\nTable: Combination of WTA-CRS with quanzation on the T5-base model.\n| Method | CoLA | MRPC | RTE | STS-B | Average | \n| :---: | :---: | :---: | :---: | :---: | :---: | \n| Full | 60.1\t| 91.5\t| 79.4\t| 90.6\t| 80.4 |\n| LoRA |60.6\t|92.2\t|80.6\t|90.7\t|81.0 |\n| WTA-CRS@0.3 |60.9\t|91.1\t|78.7\t|90.5\t|80.3 |\n| LoRA+WTA-CRS@0.3 |60     | 92\t|80.1\t|90.4\t|80.6 |\n| LoRA+WTA-CRS@0.3+Quant@8bit |60.3\t|92.06\t|81.2\t|90.4\t|81.0|\n\n\n\n\n**[W2] Discuss the difference about WTA-CRS and RAD**\n\nWTA-CRS and RAD shares the same spirit in the sense that they both trade gradient noise in return for reduced memory. However, the main difference between them lies in how they generate the noisy gradient. Specifically, WTA-CRS focues on approximating expensive matrix production operation. RAD proposes two noisy-yet-cheap gradient estimator, i.e. path sampling (sampling the computation path) and random matrix injection (apply random projection to activations). These techniques are orthogonal to each other. We will include this discussion in the updated version.\n\n"}, {"review_id": "MxzxdCzRNe", "reviewer": "Reviewer_j1mL", "summary": "In this paper, we propose a new method called WTA-CRS (Winner-Take-All Column Row Sampling) to address the main memory bottleneck issue during training, which arises from storing feature maps. To reduce memory usage during training, we sample the most likely column indices during backpropagation.\n\nFurthermore, \bthey proposed method demonstrates the ability to significantly reduce peak memory usage, by approximately up to 2.7 times, when fine-tuning downstream tasks. It also showcases the potential for higher throughput, enabling more efficient training.", "questions": "* In Figure 9, it can be observed that the throughput of WTS-CRS is lower than that of full when the batch size is small. Is this due to the overhead caused by lowering?\n* When comparing the training throughput, how does CRS differ from full in terms of throughput?\n* Could the authors include statistics for GPU utilization in their experiments? It would be helpful to analyze the causes of improved performance more thoroughly.\n* Considering that most large models are trained using multiple levels of parallelism, would it be possible to verify results for pipeline parallel, tensor parallel, etc.? Also, it is unclear from the paper whether the data parallelism used was distributed data parallelism or naïve data parallelism.", "limitations": "* As previously mentioned, it would be valuable to include additional experimental results for models that are more challenging to quantify, such as GPT-series (OPT, LLaMA). This would enhance the validity and applicability of the proposed method across a broader range of models.\n* Considering that most large-scale models are trained using multiple levels of parallelism, it is important to assess how much the proposed methods, such as pipeline parallelism and tensor parallelism, can increase throughput while taking into account overhead (such as GPU-to-GPU or node-to-node communication), memory reduction, and computational cost. Furthermore, it is not clear from the paper whether the data parallel processing used is distributed data parallelism or naive data parallelism.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The work clearly states its motivation and its solution and is easy to follow.\n2. The authors show that their method reaches comparable performance with backpropagation using the full activation when combined with LoRA.\n3. They also empirically measure throughput gains obtained by increasing batch size, which demonstrates the practical applicability of their method.", "weaknesses": "1. The paper needs a comparative analysis of other researchs, such as gradient checkpoint/recalculation and CRS, aimed at reducing activation memory during the training phase, as shown in Fig. 6 and Fig. 9.\n2. The paper should include an analysis of the overhead associated with the proposed WTS-CRS method, which involves sampling rows and columns. It is crucial to consider factors such as the computational cost of Equation 3 and any potential effects of lowering on the overall performance. Providing this analysis would enhance the clarity and completeness of the research.\n3. There is a need of analysis on the effectiveness of the proposed approach, WTS-CRS, in distributed training environments such as tensor parallelism or pipeline parallelism.\n4. It seems necessary to conduct performance evaluations on various LLMs of the GPT family, such as LLaMA and OPT.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**[W1,Q2] Compare against gradient checkpoint/recalculation and CRS**\n\nFor the comparison to CRS, from the accuracy perspective, we already compared them in Figure 8. We observe that CRS cannot maintain the accuracy. Thus, it is less applicable to fine-tuning, let alone memory saving. From the memory-saving perspective, CRS and WTA-CRS share the same implementation. Thus if the number of column-row pairs are the same, their memory saving is exactly the same.\n\nFor the comparison to gradient checkpoint, **we discussed it in Appendix A**. Here we conduct a more detailed comparison between gradient checkpointing and WTA-CRS within the huggingface backend: We set 'gradient_checkpointing = True' in the configuration and report the final memory saving.\n\nTable: Comparison between gradient checkpoint and WTA-CRS in terms of memory footprints (GB)\n| Method | T5-Base\t| T5-Large | \n| :---: | :---: | :---: | \n| FP\t| 17.66\t| 45.85\n| Grad-checkpoint\t| 13.91 (1.27x)\t| 36.5 (1.25x)\n| LoRA+WTA-CRS@0.3\t| 8.44 (2.1x)\t| 21.58 (2.1x) \n| LoRA+WTA-CRS@0.1\t| 7.30 (2.4x)\t| 18.46 (2.5×)\n\n\n**[W2,Q1] The analysis of the overhead associated with WTS-CRS**\n\nWe kindly draw your attention to **Appendix E.2**, where we already conducted an in-depth analysis. We present the table here for your convenience.\n\nThe following table provides a breakdown of the latency in our implementation: Fwd', 'Bwd', and 'F-B' represent the time of forward pass, backward pass, and the total time for both the forward and backward pass, respectively. We summarize that, under the same workload, **the current implementation of WTA-CRS may experience a roughly 20% slowdown in linear operation**. This can be attributed to the extra sampling process. \n\nAlthough we remove 70% column-row pairs, the backward time is only slightly faster than the baseline. This is mainly because the current implementation separately indexes a subset of gradient tensor before multiplying with the subsampled activations, which incurs many extra I/O [2] (also, please check Figure 13 in [2]). **Fortunately, this overhead can be sigificantly reduced with kernel fusion using Triton [3]. According to Figure 13 in [2], we expected a 2X speedup for backward pass with this Tritnon implementation [3].**\n\nTable: Latency (ms) of Forward, Backward and Forward-backward pass.\n| | Method | T5-Attention | T5-FFN | T5-Block | T5-Large |\n| :---: | :---: | :---: | :---: | :---: | :---: |\n| Fwd | Full | 8 | 10 | 17 | 1052\n| Fwd | WTA-CRS@0.3 | 22 | 16 | 37 | 2013\n| Bwd | Full | 16 | 19 | 34 | 2073\n| Bwd | WTA-CRS@0.3 | 15 | 14 | 30 | 1738\n| F-B | Full | 24 | 29 | 51 | 3125\n| F-B | WTA-CRS@0.3 | 37 | 30 | 67 | 3751\n\n\n**[W3, Q4, Limitation2] Evaluate WTS-CRS under tensor parallelism or pipeline parallelism. Also, it is unclear from the paper whether the data parallelism used was distributed data parallelism or naïve data parallelism** \n\nFor the model/tensor parallelism setting, first, up to our knowledge, it is rarely used in fine-tuning scenario [4], which is the main focus of this paper. For fine-tuning, the ideal case is to use one single GPU to tune a model as large as possible [4]. \n\nSecond, the pipeline/model parallelism requires the division of the model into smaller segments, which are distributed across multiple devices. Thus, it requires the communication of activations between consecutive model parts, potentially causing substantial overhead [5]. In this context,  WTA-CRS significantly reduces the communication volume by compressing the activation, thus reducing this overhead [5]. We leave it as future work.\n\nFor the \"data parallelism\" question, the \"data parallelism\" in this paper refers to **distributed data parallelism.**\n\n**[W4, Limitation1] Conduct performance evaluations on LLMs of GPT family, such as LLaMA and OPT.**\n\nHere we conducted additional experiments of applying WTA-CRS to the OPT model in the below table: LoRA+WTA-CRS@0.3 shows almost no drop in accuracy when compared with full training and LoRA. WTA-CRS has been applied to various transformer architectures including encoder-only (BERT), decoder-only (OPT), and encoder-decoder (T5). With this comprehensive evaluation, we believe that the experiments sufficiently demonstrate the effectiveness of WTA-CRS across diverse transformer architectures.\n\n\nTable: Experiment results on OPT-350M.\n| Method | CoLA | MRPC | RTE | STS-B | Average | \n| :---: | :---: | :---: | :---: | :---: | :---: | \n| Full\t| 49.84\t| 85.47\t| 72.56\t| 84.43\t| 73.075 \n| LoRA\t| 52.3\t| 88.36\t| 74.01\t| 87.21\t| 75.47\n| LoRA+WTA-CRS@0.3\t| 51.8\t| 88.43\t| 74.01\t| 86.61\t| 75.2125\n\n**[Q3] Statistics for GPU utilization**\n\nWTA-CRS has extra I/O (the sampling process) in return for reduced computations (FLOPs). Thus, the GPU utility of WTA-CRS is expected to lower than the standard training. Below we measure the GPU utility using torch.cuda.utilization: we presented the GPU utility during training T5-Large. Our result indicates that WTA-CRS reduces 30% GPU utilization of foward pass due to the extra sampling process. However, we note that **GPU utilization cannot reflect the wall-clock speed** : Although the GPU utility of the backward pass in full training is 100%, the wall clock time may still longer than WTA-CRS@0.3 as it requires 70% more FLOPs (workload), especially with kernel fusion implementation [3].\n\nTable: GPU utilization of training T5-Large model on NVIDIA A5000 GPU.\n| Method | Fwd | Bwd | Average |\n| :---: | :---: | :---: | :---: | \n| Full    | 75.2% | 100% | 87.6% |\n| WTA-CRS@0.3\t| 40.6% | 100% | 70.3% |\n\n[1] GACT: Activation compressed training for generic network architectures\n\n[2] Deja vu: Contextual sparsity for efficient llms at inference time\n\n[3] https://github.com/FMInference/DejaVu/blob/master/Dejavu/src/ops/triton/gather_gemv.py\n\n[4] QLoRA: Efficient Finetuning of Quantized LLMs\n\n[5] Fine-tuning Language Models over Slow Networks using Activation Quantization with Guarantees"}, {"review_id": "La7Po5nQTT", "reviewer": "Reviewer_GDNX", "summary": ">**Rebuttal:** The provided details satisfy my concerns. I think this paper should be accepted after applying the agreed changes.\n\n>**TL;DR:** **Good paper.** The proposed WTA-CRS algorithm is based on the existing CRS algorithm and is used to reduce activation memory during training. WTA-CRS achieves up to 2.7× peak memory reduction with almost no accuracy drop and enables up to 6.4× larger batch size. However, WTA-CRS comes with computational overhead, which is discussed and explore. Addressing my concerns and questions would improve my score.\n\nThe paper proposes the WTA-CRS algorithm to reduce the neural networks training activation memory, where the paper claims that activation memory is primary memory bottleneck during training. The WTA-CRS algorithm is an unbiased estimators for matrix production with reduced variance, which only requires storing the sub-sampled activations for calculating the gradient. WTA-CRS achieves up to 2.7× peak memory reduction with almost no accuracy drop and enables up to 6.4× larger batch size.\n\nThe WTA-CRS algorithm works by sampling columns and rows to create an unbiased estimation of the original GEMM for the backpropagation. The WTA-CRS algorithm does not alter the neural architecture, and therefore the inference speed is left in tact. The experimental section shows that WTA-CRS outperforms existing prior work and is compatible with existing PEFT techniques. WTA-CRS adds a computational overhead due to sampling, however, WTA-CRS enables training on much larger batch sizes, which results in a 1.2× higher training throughput.\n", "questions": "* **Q.1.** In line #43 and Figure 2 it is noted that \"storing activations (or feature maps) is the main memory bottleneck during training\". Does this hold true for all model architectures? What about LLM training where the fine-tuning batch size is usually very small?\n* **Q.2.** Why was the WTA-CRS algorithm compared to the Deterministic top-k from [1] but not to the Bernoulli-CRS from [1]? What are the key differences between WTA-CRS and Bernoulli-CRS?\n* **Q.3.** The paper proposes WTA-CRS which sacrifices computation speed at the cost of lower peak memory. There are several existing common approaches (such as gradient checkpointing and DeepSpeed) for general memory efficient training which are compatible with PEFT techniques. Why are these comparisons not explored or detailed in the main paper?\n\n[1] Adelman, Menachem, Kfir Levy, Ido Hakimi, and Mark Silberstein. \"Faster neural network training with approximate tensor operations.\" Advances in Neural Information Processing Systems 34 (2021): 27877-27889.", "limitations": "The limitations are discussed in Appendix A.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "* **S.1.** The proposed WTA-CRS algorithm tackles an important problem in existing PEFT techniques, which makes LLM PEFT training more accessible to researchers with low resources.\n* **S.2.** The paper provides a theoretical analysis on WTA-CRS.\n* **S.3.** The proposed WTA-CRS algorithm outperform existing algorithms.\n* **S.4.** An anonymized code repository is provided as part of the submission for reproduction .\n", "weaknesses": " * **W.1.** Popular existing memory efficient training techniques such as tensor rematerialization (gradient checkpointing) [2][3] and ZeRO [1] are not compared to, although some are partially discussed in Appendix A.\n* **W.2.** The experiments are conducted on single neural network architecture (T5), although the proposed technique does not seem to be confined solely to that setting.\n* **W.3.** It is common practice today to train neural networks at a lower precision (quantization), however, it is not clear whether quantization (16bit) was used. Therefore, there is insufficient proof that the combined noise of WTA-CRS and quantization would be compatible.\n\n\n**Typos.**\n* Line #62: \"Thus\" → \"Thus,\"\n* Line #240: \"mAccording\" → \"According\"\n* Line #297: \"Thus\" → \"Thus,\"\n\n[1] Ren, J., Rajbhandari, S., Aminabadi, R.Y., Ruwase, O., Yang, S., Zhang, M., Li, D. and He, Y., 2021, July. ZeRO-Offload: Democratizing Billion-Scale Model Training. In USENIX Annual Technical Conference (pp. 551-564).\n\n[2] Jain, P., Jain, A., Nrusimha, A., Gholami, A., Abbeel, P., Gonzalez, J., Keutzer, K. and Stoica, I., 2020. Checkmate: Breaking the memory wall with optimal tensor rematerialization. Proceedings of Machine Learning and Systems, 2, pp.497-511.\n\n[3] Beaumont, O., Eyraud-Dubois, L. and Shilova, A., 2021. Efficient combination of rematerialization and offloading for training dnns. Advances in Neural Information Processing Systems, 34, pp.23844-23857.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\n**[W1, Q3] Popular existing memory efficient training techniques such as tensor rematerialization (gradient checkpointing) and ZeRO are not compared to, although some are partially discussed in Appendix A.**\n\nThank you for the suggestion. We conduct a more detailed comparision between gradient checkpointing and WTA-CRS using huggingface backend: We set 'gradient_checkpointing = True' in the huggingface backend and report the final memory saving.\n\nTable: Comparison between Gradient-checkpoint and WTA-CRS in terms of the memory footprints (GB)\n| Method | T5-Base\t| T5-Large | \n| :---: | :---: | :---: | \n| FP\t| 17.66\t| 45.85\n| Grad-checkpoint\t| 13.91 (1.27x)\t| 36.5 (1.25x)\n| LoRA+WTA-CRS@0.3\t| 8.44 (2.1x)\t| 21.58 (2.1x) \n| LoRA+WTA-CRS@0.1\t| 7.30 (2.4x)\t| 18.46 (2.5×)\n\nRegarding ZeRO [1], up to our knowledge, it is mainly designed for offloading the optimizer states, which is orthogornal to activation memory saving.\n\n**[W2] The experiments are conducted on single neural network architecture (T5)**\n\n\nWe appreciate the reviewer for this thoughtful comment. We respectfully point out that we already conducted experiments on **both encoder-only architecture, such as BERT, and encoder-decoder architecture, like T5 in Table 1**.\n\nTo further respond to the reviewer's comment, we also conduct additional experiments using the decoder-only architecture, OPT. For your convenience, we summarize the results of these experiments in the below three tables: one each for BERT-Large, OPT-350M, and T5-Large.\n\nWe observed that WTA-CRS exhibits almost no drop in accuracy compared to Full training and LoRA on all three architectures. This observation strongly indicates the effectiveness of WTA-CRS, especially considering its consistent performance across diverse transformer architectures.\n\nTable: Encoder-only architecture: BERT-Large\n| Method | CoLA | MRPC | RTE | STS-B | Average | \n| :---: | :---: | :---: | :---: | :---: | :---: | \n| Full\t| 66.8\t| 89.5\t| 72.6\t| 90.2\t| 79.775\n| LoRA\t| 65.9\t| 90.8\t| 71.3\t| 90.3\t| 79.575\n| LoRA+WTA-CRS@0.3\t| 66\t| 89.7\t| 72.4\t| 89.7\t| 79.45\n\nTable: Decoder-only architecture: OPT-350M\n| Method | CoLA | MRPC | RTE | STS-B | Average | \n| :---: | :---: | :---: | :---: | :---: | :---: | \n| Full\t| 49.84\t| 85.47\t| 72.56\t| 84.43\t| 73.075 \n| LoRA\t| 52.3\t| 88.36\t| 74.01\t| 87.21\t| 75.47\n| LoRA+WTA-CRS@0.3\t| 51.8\t| 88.43\t| 74.01\t| 86.61\t| 75.2125\n\nTable: Encoder-Decoder architecture: T5-Large\n| Method | CoLA | MRPC | RTE | STS-B | Average | \n| :---: | :---: | :---: | :---: | :---: | :---: | \n| Full\t| 61.3\t| 93.4\t| 85.3\t| 91.8\t| 82.95 | \n| LoRA\t| 63.3\t| 93.5\t| 84.2\t| 91.7\t| 83.175 | \n| LoRA+WTA-CRS@0.3\t| 62.9\t| 93.6\t| 83.9\t| 91.3\t| 82.925\n\n\n\n**[W3] It is not clear whether quantization (16bit) was used. Therefore, there is insufficient proof that the combined noise of WTA-CRS and quantization would be compatible.**\n\nWe sincerely appreciate the thoughtful comment provided by the reviewer. We agree with the reviewer's suggestion that including the experiment results of WTA-CRS with bfloat16 quantization is essential to demonstrate its effectiveness under different settings.\n\nTo address this concern, we have conducted additional experiments of WTA-CRS@0.3 on the T5-Base model with bfloat16 quantization applied to both the weight and activation map during training. The results of this experiment are presented in the following table. It is revealed that WTA-CRS@0.3 with bfloat16 quantization shows almost no drop in accuracy when compared with Full training and LoRA. This result explicitly demonstrates the effectiveness of WTA-CRS even combined with the bfloat16 quantization, indicating its robustness against the underflow noise caused by bfloat16 quantization.\n\nTable: Accuracy of WTA-CRS@0.3 on T5-Base with bfloat16 quantization.\n| Method | CoLA | MRPC | RTE | STS-B | Average | \n| :---: | :---: | :---: | :---: | :---: | :---: | \n| FP32\t| 60.1\t| 91.5\t| 79.4\t| 90.6\t| 80.4\n| LoRA-FP32\t| 60.6\t| 92.2\t| 80.6\t| 90.7\t| 81.0\n| LoRA+WTA-CRS@0.3-FP32\t| 60\t| 92\t| 80.1\t| 90.4\t| 80.6\n| LoRA+WTA-CRS@0.3-BF16\t| 60.3\t| 92.4\t| 80.1\t| 90.19\t| 80.7\n\n\n**[Q1] Are the activation memory still the bottleneck for LLM training?**\n\nThis is a great question. The short answer is the bottleneck depends on the fine-tuning setting. In data parallel/single GPU training setting, we need to store the model and optimizer states in GPU memory, and the left space is for holding activations. Thus, when the size of LLM goes beyond a certain threshold, the model weight/optmizer state must become the memory bottleneck. This is also why LLM fine-tuning often comes with a small batch size, resulting in a low GPU utility. However, when the size of LLM becomes too large such that it and its optimizer cannot be hold into one single GPU, we must tune it with pipeline/tensor/model parallelism. In such setting, it requires the division of the model into smaller segments, which are then distributed across multiple devices. In this case, each GPU only hold a small part of models, and the space left for activations become much larger. In this case, if we enlarge the sequential length and/or batch size, the activation is still the bottleneck.\n\n**[Q2] Why was the WTA-CRS algorithm compared to the Deterministic top-k from [1] but not to the Bernoulli-CRS from [1]?**\n\nWe thank the reviewer for this comment. We compare our algorithm to Deterministic top-k mainly because Deterministic Top-K works better than Bernoulli-CRS  in the context of neural network (Theorem 2 in [2], Figure 1b and Figure 3 in [2]). Thus we follow [2] to compare against top-k CRS instead of Bernoulli-CRS.\n\n**[Typos] Line #62: \"Thus\" → \"Thus,\"; Line #240: \"mAccording\" → \"According\"; Line #297: \"Thus\" → \"Thus,\".**\n\nWe appreciate the reviewer for the thoughful comments. We will fix these typos in our camera-ready version.\n\n[1] ZeRO: Memory Optimizations Toward Training Trillion Parameter Models\n\n[2] Faster Neural Network Training with Approximate Tensor Operations"}]

Author rebuttal:
We thank all the reviewers for their constructive comments and helpful feedback. We are encouraged to find that they have found our contributions to be technically solid (ky3t, cMiu, GDNX), timely and relevant for LLM research (cMiu, GDNX), mathematically solid (cMiu, GDNX), and easy-to-follow (cMiu, j1mL, GDNX).

We have additionally performed experiments to address some of the evaluation concerns. Please find below our detailed response to the questions and any concern raised by the reviewers. We will incorporate all these comments and comprehensive experimental evaluations into the revised manuscript. We are grateful to the reviewers for all the suggestions to improve our work.

Best regards,

Authors

## Summary of Rebuttal

We thank all the reviewers for their constructive comments and helpful feedback. We value their comments sincerely, and do our best to address the concerns. During the rebuttal, we provide the following new supplementary results and analysis:

- (cMiu) We have conducted an additional experiment that combines WTA-CRS with activation map quantization. [[Redirection]](https://openreview.net/forum?id=SquMNyrk1O&noteId=S24PVc9nC9)
- (cMiu, GDNX) We provide a detailed discussion on the distinctions between WTA-CRS and RAD [[Redirection]](https://openreview.net/forum?id=SquMNyrk1O&noteId=S24PVc9nC9), as well as ZeRO [[Redirection]](https://openreview.net/forum?id=SquMNyrk1O&noteId=3w86hHloPJ). 
- (j1mL) We present a technical comparison of WTA-CRS with CRS and Gradient-checkpointing, focusing on accuracy and memory cost.  [[Redirection]](https://openreview.net/forum?id=SquMNyrk1O&noteId=C8D3lWSbd2)
- (j1mL) We emphasize the importance of Appendix E.2 in our paper, as it addresses the overhead analysis of WTA-CRS, as requested by the reviewer. [[Redirection]](https://openreview.net/forum?id=SquMNyrk1O&noteId=C8D3lWSbd2)
- (j1mL) We conduct an in-depth analysis of the performance of WTA-CRS in distributed training environments. [[Redirection]](https://openreview.net/forum?id=SquMNyrk1O&noteId=C8D3lWSbd2)
- (j1mL) We provide statistics for GPU utilization to offer further insights into the efficiency of WTA-CRS. [[Redirection]](https://openreview.net/forum?id=SquMNyrk1O&noteId=C8D3lWSbd2)
- (j1mL, GDNX) We include an additional architecture, OPT-350M, which is a decoder-only transformer, in our experiments, to demonstrate the effectiveness of WTA-CRS. [[Redirection]](https://openreview.net/forum?id=SquMNyrk1O&noteId=3w86hHloPJ)
- (GDNX) We have conducted an additional experiment involving the deployment of WTA-CRS in bfloat16 fine-tuning. [[Redirection]](https://openreview.net/forum?id=SquMNyrk1O&noteId=3w86hHloPJ)

Meta-review: The paper presents an interesting novel use use of approximate matrix multiplication to parameter-efficient tuning of LLMs. The approach is theoretically-motivated, and brings significant practical memory improvements. 
Given that the reviewer feedback is (almost) uniformly positive, and the authors have committed to several small improvements over the discussion period, providing significant additional data, I am happy to recommend acceptance.
Decision: Accept (poster)