# Winner-Take-All Column Row Sampling for Memory Efficient Adaptation of
Language Model

Zirui Liu Equal contribution. The order of authors is determined by flipping a
coin. Department of Computer Science, Rice University  Guanchu Wang*
Department of Computer Science, Rice University  Shaochen Zhong  Department of
Computer Science, Rice University  Zhaozhuo Xu  Department of Computer
Science, Rice University  Daochen Zha  Department of Computer Science, Rice
University  Ruixiang Tang  Department of Computer Science, Rice University
Zhimeng Jiang  Department of Computer Science, Texas A&M University  Kaixiong
Zhou  Department of Computer Science, Rice University  Vipin Chaudhary
Department of Computer and Data Sciences, Case Western Reserve University
Shuai Xu  Department of Computer and Data Sciences, Case Western Reserve
University  Xia Hu  Department of Computer Science, Rice University

###### Abstract

With the rapid growth in model size, 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, neural networks are usually trained using stochastic
gradient descent. We argue that in stochastic optimization, models can handle
noisy gradients as long as the gradient estimator is unbiased with reasonable
variance. Following this motivation, we propose a new family of unbiased
estimators called WTA-CRS , for matrix production with reduced variance, which
only requires storing the sub-sampled activations for calculating the
gradient. Our work provides both theoretical and experimental evidence that,
in the context of tuning transformers, our proposed estimators exhibit lower
variance compared to existing ones. By replacing the linear operation with our
approximated one in transformers, we can achieve up to 2.7$\times$ peak memory
reduction with almost no accuracy drop and enables up to $6.4\times$ larger
batch size. Under the same hardware, WTA-CRS enables better down-streaming
task performance by applying larger models and/or faster training speed with
larger batch sizes.

##  1 Introduction

Pre-trained language models (LMs) with transformer architecture have achieved
remarkable success in numerous natural language processing (NLP) tasks
(Vaswani et al., 2017; Devlin et al., 2018; Chuang et al., 2023; Raffel et
al., 2020b; Brown et al., 2020; Yang et al., 2023; Zha et al., 2023).
Specifically, these models are trained on vast text corpora to acquire
general-purpose representations, which are then adapted to a specific task by
fine-tuning on task-specific data. In recent studies, it has been convincingly
demonstrated that significantly increasing the number of parameters in pre-
trained LMs leads to remarkable improvements in performance (Kaplan et al.,
2020). As a result, there is now an urgent necessity to effectively adapt
these models, equipped with billion-scale parameters, to a wide range of
tasks.

However, a significant disparity exists between the memory requirements of
pre-trained LMs and the capacity of current hardware, particularly GPUs. For
example, even a GPU with 24GB memory cannot accommodate the fine-tuning
process of the T5-3B model (Raffel et al., 2020b) with batch size one, which
boasts three billion parameters. Without additional techniques, attempting to
fine-tune billion-scale LMs on a single GPU is impossible. Although model-
parallel fine-tuning is feasible, the majority of the time, we cannot bear the
expense of acquiring multiple GPUs or the communication overhead involved. To
ensure the smooth deployment of language models during the fine-tuning
process, it is crucial to adapt them for operation on a single GPU.

To address this issue, several parameter-efficient tuning methods are proposed
(Lester et al., 2021; Sung et al., 2022; Li and Liang, 2021; Zaken et al.,
2021; Hu et al., 2021; Karimi Mahabadi et al., 2021; Houlsby et al., 2019).
Specifically, adapters (Houlsby et al., 2019; Karimi Mahabadi et al., 2021)
insert a small module into the transformer blocks and only update it while
keeping other parameters fixed. Similarly, prompt tuning (Lester et al., 2021)
introduces a small vector that is concatenated with the input embeddings and
updated during the tuning process. LoRA (Hu et al., 2021) 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 (Kingma and Ba, 2014; Hu et al., 2021).
Although the optimizer states contribute to the memory footprint, _storing
activations (or feature maps) is the main memory bottleneck during training_
(often $>70\%$) (Chen et al., 2021c; Jain et al., 2020; Kirisame et al., 2020;
Andoorveedu et al., 2022b). Thus, parameter-efficient methods often do not
reduce memory usage by much (Sung et al., 2022; Andoorveedu et al., 2022b).

Figure 1: Accuracy-memory trade-off of WTA-CRS and other memory-efficient
tuning methods. Unless specially stated, we use the T5-Large in the figure.

In parallel, we can reduce the main memory bottleneck by reducing the
activation storage in fine-tuning. Since transformer-based models are mainly
built based on the linear layer, a less-explored direction is to replace the
expensive matrix multiplication operation with its memory-efficient
estimations using column-row sampling (CRS) (Adelman et al., 2021; Drineas et
al., 2006). The key idea of CRS is to sub-sample tensors onto low-dimensional
spaces and perform the original operations here. Specifically, for the linear
operation between two matrices $\mathbf{A}\in\mathbb{R}^{n\times m}$ and
$\mathbf{B}\in\mathbb{R}^{m\times q}$ (in the context of machine learning,
$\mathbf{A}$ is often activations), we first sample $k$ ($k<m$) column-row
pairs according to a pre-defined distribution. Then we obtain
$\mathbf{A}^{\prime}\in\mathbb{R}^{n\times k}$ and
$\mathbf{B}^{\prime}\in\mathbb{R}^{k\times q}$ ($k<m$) by picking $k$ columns
of $\mathbf{A}$ and the corresponding rows of $\mathbf{B}$ according to the
sampled column-row pairs (Drineas et al., 2006). Finally, we estimate
$\mathbf{A}\mathbf{B}\approx\mathbf{A^{\prime}}\mathbf{B^{\prime}}$. In this
way, we only need to store the sub-matrix $\mathbf{A}^{\prime}$ and
$\mathbf{B}^{\prime}$ in GPU memory to perform the computation. Moreover,
transformer-based models training/tuning are performed with the first-order
stochastic optimizer, e.g., Adam (Kingma and Ba, 2014). In stochastic
optimization, models can work with noisy gradients, _as long as the gradient
estimator is unbiased and has a reasonable variance._ In view of such, we ask:
why spend resources on obtaining exact gradients when we are using stochastic
optimization? Motivated by this, we focus on obtaining unbiased gradients
cheaply with approximated matrix multiplication.

The approximation method reduces the memory usage at the cost of giving
outputs with variance. Thus there naturally exists an accuracy-memory trade-
off. The main challenge is how to integrate the approximated matrix
multiplication into transformer with minimal gradient variance. In this paper,
we propose a new family of unbiased estimator for matrix multiplication with
reduced variance, dubbed Winner-Take-All Column-Row Sampling (WTA-CRS ).
Compared to CRS, WTA-CRS reduces the variance of an estimator by focusing more
on high-probability regions of the sampling distribution. Moreover, WTA-CRS
can serve as a drop-in replacement for the linear operation in transformers,
providing an unbiased weight gradient with reduced memory usage. As shown in
Figure 1, our method achieves better accuracy-memory trade-off than state-of-
the-art memory-efficient tuning methods, e.g., LST (Sung et al., 2022) and
LoRA (Hu et al., 2021). Moreover, since WTA-CRS executed at the operation
level, it is orthogonal to most of the existing parameter-efficient tuning
methods. Our contributions are highlighted as follows:

  * •

We design a new family of unbiased estimator for matrix multiplication with
reduced variance. We theoretically and experimentally verify that it has
smaller variance than the established one under the context of tuning
transformer.

  * •

By replacing the linear operation with WTA-CRS in transformers, we can achieve
up to 2.7$\times$ peak memory reduction with almost no accuracy drop, and
enables up to $6.4\times$ larger batch size. As shown in Figure 1, WTA-CRS
stands out as an exceptional solution capable of fine-tuning T5-3B using a
mere 40GB GPU memory budget, with three billion parameters. Thus, we achieve
remarkable advancements in the adaptation of LMs on downstream tasks.

  * •

We implement WTA-CRS as a ready-to-use extension for Pytorch with an easy-to-
use API that can also be combined with other memory-saving techniques.

##  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

Figure 2: The GPU memory usage breakdown for fine-tuning T5 (Wang et al.,
2018a), where the batch size $B$ is 64 and sequential length $S$ is 128 or
256.

In each training step of backpropagation, it has exactly two phases, i.e., one
forward phase and one backward phase. Transformer-based models are mainly
built based on the linear operation, which can be written as:

$\displaystyle\texttt{Forward Pass}~{}~{}~{}~{}~{}~{}~{}~{}{\mathbf{Z}}$
$\displaystyle=\texttt{GEMM}({\color[rgb]{1,0,0}{\mathbf{H}}},{\mathbf{W}}),$
(1a) $\displaystyle\texttt{Backward
Pass}~{}~{}~{}~{}~{}~{}~{}~{}\nabla{\mathbf{H}}$
$\displaystyle=\texttt{GEMM}(\nabla{\mathbf{Z}},{\mathbf{W}}^{\top}),$ (1b)
$\displaystyle\nabla{\mathbf{W}}$
$\displaystyle=\texttt{GEMM}({\color[rgb]{1,0,0}{\mathbf{H}}^{\top}},\nabla{\mathbf{Z}}),$
(1c)

where $\texttt{GEMM}(\cdot,\cdot)$ is the General Matrix Multiplication
operation, ${\mathbf{H}}$ and ${\mathbf{Z}}$ are the activation (or input
feature maps) and output feature maps, respectively. ${\mathbf{W}}$ is the
weight of the linear layer. $\nabla{\mathbf{H}}$, $\nabla{\mathbf{W}}$, and
$\nabla{\mathbf{Z}}$ are the gradient of ${\mathbf{H}}$, ${\mathbf{W}}$, and
${\mathbf{Z}}$, respectively. From Equation (1c), activations ${\mathbf{H}}$
are used in the backward phase. In commonly used deep learning framework
(Abadi et al., 2016; Paszke et al., 2019), it requires storing ${\mathbf{H}}$
in GPU memory during the forward pass, for calculating the weight gradient
$\nabla{\mathbf{W}}$ in the backward pass.

Previous works show that although the model parameters contribute to the
memory footprint, activations (e.g., storing ${\mathbf{H}}$) are the main
memory bottleneck during training (Chen et al., 2021c; Jain et al., 2020;
Kirisame et al., 2020; Andoorveedu et al., 2022b). To get a sense of the
scale, we show in Figure 2 that for popular transformer models like T5,
activations may take roughly $73\sim 88\%$ of the total memory, depending on
the batch size $B$ and sequential length $S$.

###  2.2 Approximated GEMM With Sampling

Let ${\mathbf{X}}\in\mathbb{R}^{n\times m}$,
${\mathbf{Y}}\in\mathbb{R}^{m\times q}$ be two matrices. The goal is to
efficiently estimate the matrix production ${\mathbf{X}}{\mathbf{Y}}$.
Singular Value Decomposition (SVD) outputs provably optimal low-rank
estimation of ${\mathbf{X}}{\mathbf{Y}}$ (Adelman et al., 2021). However, SVD
is almost as expensive as matrix production itself. Instead, the sampling
algorithm is proposed to approximate the matrix product
${\mathbf{X}}{\mathbf{Y}}$ by sampling $k$ columns of ${\mathbf{X}}$ and
corresponding rows of ${\mathbf{Y}}$ to form smaller matrices, which are then
multiplied as usual (Drineas et al., 2006; Drineas and Kannan, 2001):

$\displaystyle\texttt{GEMM}({\mathbf{X}},{\mathbf{Y}})$
$\displaystyle=\sum_{i=1}^{m}{\mathbf{X}}_{:,i}{\mathbf{Y}}_{i,:}\approx\sum_{t=1}^{k}\frac{1}{kp_{i_{t}}}{\mathbf{X}}_{:,i_{t}}{\mathbf{Y}}_{i_{t},:}={\mathbf{X}}^{\prime}{\mathbf{Y}}^{\prime},$
(2)

where ${\mathbf{X}}_{:,i}\in\mathbb{R}^{n\times 1}$ and
${\mathbf{Y}}_{i,:}\in\mathbb{R}^{1\times q}$ are the $i^{\mathrm{th}}$ column
and row of ${\mathbf{X}}$ and ${\mathbf{Y}}$, respectively. In this paper, we
call $({\mathbf{X}}_{:,i},{\mathbf{Y}}_{i,:})$ the $i^{\mathrm{th}}$ column-
row pair. $k$ is the number of sampled pairs ($1\leq k\leq m$).
$\mathcal{P}=\\{p_{i}\\}_{i=1}^{m}$ is a probability distribution over the
column-row pairs. $i_{t}\in\\{1,\cdots m\\}$ is the index of the sampled
column-row pair at the $t^{\mathrm{th}}$ trial. $s_{t}$ is the scale factor.
${\mathbf{X}}^{\prime}\in\mathbb{R}^{n\times k}$ and
${\mathbf{Y}}^{\prime}\in\mathbb{R}^{k\times q}$ are the normalized sub-
matrices sliced according to the sampled column-row pairs.

Existing work (Drineas et al., 2006) shows
${\mathbf{X}}^{\prime}{\mathbf{Y}}^{\prime}$ is an unbiased estimation of
${\mathbf{X}}{\mathbf{Y}}$, i.e.,
$\mathbb{E}[{\mathbf{X}}^{\prime}{\mathbf{Y}}^{\prime}]={\mathbf{X}}{\mathbf{Y}}$.
Furthermore, the approximation error
$\mathbb{E}[||{\mathbf{X}}{\mathbf{Y}}-{\mathbf{X}}^{\prime}{\mathbf{Y}}^{\prime}||_{F}]$
is minimized when the probabilities $\\{p_{i}\\}_{i=1}^{m}$ are proportional
to the product of the column-row Euclidean norms (Drineas et al., 2006) (Proof
in Appendix C):

$p_{i}=\frac{||{\mathbf{X}}_{:,i}||_{2}\
||{\mathbf{Y}}_{i,:}||_{2}}{\sum_{j=1}^{m}||{\mathbf{X}}_{:,j}||_{2}\
||{\mathbf{Y}}_{j,:}||_{2}}.$ (3)

As we analyzed in Section 2.1, storing the activation ${\mathbf{H}}$ is the
major memory bottleneck.  If we can replace
$\texttt{GEMM}({\mathbf{H}}^{\top},\nabla{\mathbf{Z}})$ in Equation (1c) with
${\mathbf{H}}^{\prime\top}\nabla{\mathbf{Z}}^{\prime}$ following the paradigm
of Equation (2), then we only need ${\mathbf{H}}^{\prime}$ instead of
${\mathbf{H}}$ in GPU memory to compute the gradient, which significantly
decreases the memory usage of activations. This estimation linearly reduces
the memory complexity from $\mathcal{O}(nm)$ to $\mathcal{O}(nk)$. Also, the
total number of floating point operations (FLOPs) is reduced as well since the
computation is executed on two smaller matrices. _For the ease of
illustration, in this paper we call the distribution in Equation ( 3) the
column-row index distribution_. In the next section, we explore how to reduce
memory usage via sampling-based matrix multiplication.

##  3 Methodology

In recent years, we have observed that deep neural network training can be
performed almost entirely with first-order _stochastic optimization_ (Kingma
and Ba, 2014). Thus intuitively, _in stochastic optimization we can reduce the
resources spent on obtaining gradients, as long as the estimated gradient is
unbiased with reasonable variance_ (Chmiel et al., 2023, 2021; Oktay et al.,
2020). Following this motivation, we first design a new unbiased estimator for
matrix multiplication with reduced variance compared to the one in Equation
(2) (Section 3.1 ). Then we introduce how to replace the GEMM in Transformer
with its approximated version to reduce the memory usage (Section 3.2).

###  3.1 Winner-Take-All Column-Row Sampling: A New Unbiased Estimator for
GEMM

In this section, we mathematically design a new unbiased estimator for GEMM
with reduced variance called WTA-CRS (Winner-Take-All Column-Row Sampling).
Following the notation in Section 2.2, let ${\mathbf{X}}\in\mathbb{R}^{n\times
m}$, ${\mathbf{Y}}\in\mathbb{R}^{m\times q}$ be two matrices.
$\mathcal{P}=\\{p_{i}\\}_{i=1}^{m}$ is the column-row index distribution in
Equation (3)111Here we note that the theoretical analysis in this section can
be applied to any probability distribution, not only limited to the one in
Equation (3).. We first define the variable $f(i)$ as

$f(i)=\frac{{\mathbf{X}}_{:i}{\mathbf{Y}}_{i:}}{p_{i}},$

$f(i)$ is an unbiased estimation for the matrix production between
${\mathbf{X}}$ and ${\mathbf{Y}}$. To see this,

$\displaystyle\mathbb{E}_{j\sim\mathcal{P}}[f(j)]=\sum_{i=1}^{m}p_{i}\frac{{\mathbf{X}}_{:,i}{\mathbf{Y}}_{i:}}{p_{i}}={\mathbf{X}}{\mathbf{Y}}.$

We note that the prior approximated matrix multiplication in Equation (2) is
the direct extension of $f(i)$ by taking the average of
$\\{f(i_{t})\\}_{t=1}^{k}$ among $k$ independent random trials to reduce the
variance. Here we explore an alternative approach to reduce the variance of
$f(i)$ beyond simple averaging. Our core idea is to partition the column-row
index distribution $\mathcal{P}=\\{p_{i}\\}_{i=1}^{m}$ into two complementary
regions based on the probability mass: a high-probability region
$\mathcal{P}^{\mathcal{C}}$ and a low-probability region
$\mathcal{P}^{\mathcal{D}\backslash\mathcal{C}}$, where
$\mathcal{D}=\\{1,\cdots,m\\}$ is the whole set and $\mathcal{C}$ is the set
of the column-row index with the largest probability. Let $\mathcal{C}$ be the
set of column-row pair indices associated with $|\mathcal{C}|$ largest
$p_{i}$. We define WTA-CRS estimator for ${\mathbf{X}}{\mathbf{Y}}$ as
follows:

$\displaystyle\mathbb{E}_{j\sim\mathcal{P}^{\mathcal{D}\backslash\mathcal{C}}}\Big{[}\sum_{c\in\mathcal{C}}f(c)p_{c}+(1-\sum_{c\in\mathcal{C}}p_{c})f(j)\Big{]}.$
(4)

We note that the random variable in Equation (4) is the column-row pair index
$j$, and is only sampled from $\mathcal{D}\backslash\mathcal{C}$. The
estimator defined in Equation (4) contains two parts. The first part
$\sum_{c\in\mathcal{C}}f(c)p_{c}$ has no relationship with the random variable
$j$ and is summed deterministically. The second part $f(j)$ is sampled
stocastically, but scaled by the factor $(1-\sum_{c\in\mathcal{C}}p_{c})$.
When $\mathcal{P}=\\{p_{i}\\}_{i=1}^{m}$ is concentrated on a small number of
atoms, the scaling factor $(1-\sum_{c\in\mathcal{C}}p_{c})$ for the stochastic
term should be small. Therefore, we intuitively expect the estimator to have a
small variance in this case due to a small scaling factor. In this way, we
reduce the variance of an estimator by focusing more on high-probability
regions of the distribution (winner-take-all). Below we formalize this
intuition by showing the statistical property of our estimator regarding the
bias and variance, respectively.

######  Theorem 1 (Proof in Appendix C.2).

The estimator defined in Equation (4) is an unbiased estimator for matrix
production ${\mathbf{X}}{\mathbf{Y}}$, i.e,
$\mathbb{E}_{j\sim\mathcal{P}^{\mathcal{D}\backslash\mathcal{C}}}[\sum_{c\in\mathcal{C}}f(c)p_{c}+(1-\sum_{c\in\mathcal{C}}p_{c})f(j)]={\mathbf{X}}{\mathbf{Y}}$.

Theorem 1 states that our proposed estimator in Equation (4) is unbiased.
Below we compare our proposed estimator to the CRS estimator in Equation (2)
in terms of the variance. Suppose we have the budget of only utilizing $k$
column-row pairs for approximating the matrix production. From the
implementation perspective, the estimator defined in Equation (2) estimates
$\texttt{GEMM}({\mathbf{X}},{\mathbf{Y}})$ as:

$\displaystyle\texttt{(CRS)}~{}~{}~{}~{}~{}~{}g({\mathbf{X}},{\mathbf{Y}})=\frac{1}{k}\sum_{t=1}^{k}f(i_{t}),~{}~{}~{}i_{1},\cdots,i_{k}\overset{\text{i.i.d}}{\sim}\mathcal{P}.$
(5)

Our estimator defined in Equation (4) splits the budget $k$ into two parts.
Namely, the first part explicitly sums the expectation terms for the largest
probability group $\mathcal{C}$ ($|\mathcal{C}|<k$), while stochastically
average $k-|\mathcal{C}|$ samples drawn from
$\mathcal{D}\backslash\mathcal{C}$ to estimate the remaining terms, up to
scale:

$\displaystyle(\texttt{WTA-
CRS})~{}~{}~{}\hat{g}({\mathbf{X}},{\mathbf{Y}})=\sum_{c\in\mathcal{C}}f(c)p(c)+\frac{1-\sum_{c\in\mathcal{C}}p_{c}}{k-|\mathcal{C}|}\sum_{j=1}^{k-|\mathcal{C}|}f(j),~{}~{}~{}i_{1},\cdots,i_{k-|\mathcal{C}|}\overset{\text{i.i.d}}{\sim}\mathcal{P}^{\mathcal{D}\backslash\mathcal{C}}.$
(6)

######  Theorem 2 (Proof in Appendix C.3).

Suppose the total budget of column-row pairs is $k$. If $\mathcal{C}$
satisfies

$\sum_{c\in\mathcal{C}}p_{c}>\frac{|\mathcal{C}|}{k},$ (7)

then we have
$\mathrm{Var}[\hat{g}({\mathbf{X}},{\mathbf{Y}})]<\mathrm{Var}[g({\mathbf{X}},{\mathbf{Y}})]$.
Moreover, $\mathrm{Var}[\hat{g}({\mathbf{X}},{\mathbf{Y}})]$ is minimized when
$|\mathcal{C}|=\min_{|\mathcal{C}|\in\\{0,\cdots,k\\}}\frac{1-\sum_{c\in\mathcal{C}}p_{c}}{k-|\mathcal{C}|}$.

Both the left- and right-hand sides of Equation (7) depend on the size of the
highest probability group $|\mathcal{C}|$, which controls the number of high
probability column-row pairs that are directly added without sampling. Below
we experimentally investigate whether Equation (7) holds under the context of
fine-tuning the transformer-based model with varying $|\mathcal{C}|$.

Figure 3: The probability mass $\sum_{c\in\mathcal{C}}p_{c}$ versus
$\frac{|\mathcal{C}|}{k}$ in Equation (7) at $k=0.3|\mathcal{D}|$. Here we
visualize the column-row index distribution of query/key/value projection
layer in the T5-base model, which is fine-tuned on RTE dataset. More similar
results can be found in Appendix E.1.

Experimental analysis. As shown in Figure 3, we visualize the two terms in
Equation (3) for the column-row index distribution of query, key, and value
projection in the self-attention module, respectively (Vaswani et al., 2017).
Specifically, we fix the total column-row pair budget $k=0.3|\mathcal{D}|$ and
change the size of the highest probability group $|\mathcal{C}|$ from 0 to
$k$. We conclude that Equation (7) holds for most of the layers when fine-
tuning transformers. Thus, we expect our WTA-CRS has better performance than
CRS for adapting transformer-based models, which is later experimentally
verified in Section 5.

###  3.2 Compress GEMM in Transformers with WTA-CRS

Previous work has shown that unbiasedness of the estimated gradient is crucial
for the proper convergence of stochastic gradient descent (Chmiel et al.,
2023, 2021; Chen et al., 2021c; Liu et al., 2022a). As shown in Section 2.1,
we have three GEMM in the linear layer. Below we investigate how to replace
GEMM with its approximated version in a way that the estimated gradient is
unbiased.

Unbiasedness. Previous work has shown that to ensure the unbiasedness of the
gradient, the approximation can only be applied during the backward pass (Chen
et al., 2021c; Liu et al., 2022b; Adelman et al., 2021). The rationale behind
this conclusion is that we have $\mathbb{E}[f(x)]\neq f(\mathbb{E}[x])$ for
any non-linear function $f(\cdot)$, e.g.,
$\mathbb{E}[x^{2}]\neq\mathbb{E}^{2}[x]$. Thus if we replace the forward GEMM
in Equation (1a), even when the approximation method gives an unbiased
estimation, i.e.,
$\mathbb{E}[\hat{g}({\mathbf{H}},{\mathbf{W}})]={\mathbf{H}}{\mathbf{W}}={\mathbf{Z}}$,
the output activations (e.g., $\textrm{GeLU}({\mathbf{Z}})$) are still biased
since the activation function is non-linear, namely,

$\textrm{GeLU}(\hat{g}({\mathbf{H}},{\mathbf{W}}))=\textrm{GeLU}(\mathbb{E}[{\mathbf{Z}}])\neq\mathbb{E}[\textrm{GeLU}({\mathbf{Z}})].$

To ensure the unbiasness of gradient and reduce the memory usage of storing
${\mathbf{H}}$, as shown in the example of Figure 5, we only replace GEMM in
the backward pass with its approximation (e.g., Equation (1c)), while leaving
the forward one unchanged (e.g., Equation (1a)). We show in Appendix B that
the estimated weight gradient is unbiased in this case.

Figure 4:  The diagram of a single Transformer block. The shape of activations
is annotated, where $B,S,D_{\text{model}}$, $N_{\text{head}}$, and
$D_{\text{head}}$ are the batch size, sequence length, hidden size, number of
attention heads, and head dimension, respectively. WTA-CRS  can be applied
to the operators in green; the activation maps of operators in blue can be
losslessly compressed; and those in gray are not compressed in this paper. The
idea of this figure is inspired by (Andoorveedu et al., 2022a). Figure 5: The
illustration of how to deploy WTA-CRS to linear layers. We only replace GEMM
in Equation (1c) with its approximated version using WTA-CRS . The pseudocode
is given in Appendix D Algorithm 1.

#### Implementation.

Here we present how we implement WTA-CRS in Equation (6) in practice. For the
linear layer, as we analyzed, we only replace GEMM in Equation (1c) with its
approximated version. In this case, ${\mathbf{X}}$ and ${\mathbf{Y}}$ in
Equation (6) are activation ${\mathbf{H}}^{\top}$ and output gradient
$\nabla{\mathbf{Z}}$, respectively. Given the total column-row pair budget
$k$, the first step is to build the deterministic index set $\mathcal{C}$,
where each element is summed explicitly without sampling. Note that
$\mathcal{C}$ is a set of indices with the highest probabilities in Equation
(3). Thus, to build $\mathcal{C}$, we only need to determine its size, denoted
as $|\mathcal{C}|$, which minimizes the variance of the estimator. As Theorem
2 suggested, we set
$|\mathcal{C}|=\min_{|\mathcal{C}|\in\\{0,\cdots,k\\}}\frac{1-\sum_{c\in\mathcal{C}}p_{c}}{k-|\mathcal{C}|}$.
The second step is to sample $k-|\mathcal{C}|$ column-row indices from the
remaining distribution $\mathcal{P}^{\mathcal{D}\backslash\mathcal{C}}$ to
obtain the set $\mathcal{C}_{\text{stoc}}$, where
$|\mathcal{C}_{\text{stoc}}|=k-|\mathcal{C}|$. The third step is to build sub-
sampled ${\mathbf{H}}^{\prime}$ with only rows from
$\mathcal{C}\cup\mathcal{C}_{\text{stoc}}$. Note that for rows in
${\mathbf{H}}^{\prime}$ from $\mathcal{C}_{\text{stoc}}$, we need to normalize
it by $\frac{1-\sum_{c\in\mathcal{C}}p_{c}}{k-|\mathcal{C}|}$ according to
Equation (6). We illustrate the above process in Figure 5. The pseudocode to
Appendix D Algorithm 1.

#### Scope.

Here we show which operation can be replaced with its approximation version.
As shown in Figure 4, the transformer is mainly consisted of linear layer,
TensorMul, and other operations (e.g., GeLU, Dropout, LayerNorm). TensorMul in
Figure 4 refers to the multiplication between two four-dimensional tensors.
Our WTA-CRS  _can be applied to Linear-Q, -K, -V, -O, -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.

##  4 Related Work and Discussion

Due to the page limit, we discuss the related work on approximated matrix
multiplication and activation compression. Other related topics, e.g.,
parameter-efficient fine-tuning and gradient checkpointing, can be found in
Appendix A. We also discuss the limitation and potential negative social
impact in Appendix A.

Approximated Matrix Multiplication. In the context of neural networks,
approximated matrix multiplication methods can be broadly categorized into two
main groups: (1) Butterfly-based methods (Chen et al., 2021a; Dao et al.,
2022) 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 (Drineas et al., 2006; Adelman et al., 2021; Liu et al., 2022b)
select important rows and columns from the input matrices and perform the
multiplication on the sampled matrix. Our work is closely aligned with this
second research line. (Adelman et al., 2021; Liu et al., 2022b) 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, (Adelman et al., 2021)
deterministically selects column-row pairs without scaling, whereas our
estimator divides the column-row pairs into a deterministic component and a
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 8).

Activation Quantization. The activation quantization methods focus on
quantizing the activation into low numerical precision numbers, e.g., 8-bit
integers (Chen et al., 2021c; Liu et al., 2022a, c; Wang et al., 2022b, a).
Here we discuss the difference between these two works in terms of compression
ratios. According to Table 5 in (Liu et al., 2022a), when GACT is combined
with Swapping, i.e., offloading quantized activation to main memory, it
achieves a peak memory usage compression rate of $1.73\times$ for Bert-Large.
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.

##  5 Experiments

In this section, we design experiments to answer the following research
questions: RQ1: 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, we adopt GLUE benchmark (Wang et al.,
2018b) to evaluate the effectiveness of different methods, including the CoLA,
SST-2, MRPC, QQP, MNLI, QNLI, RTE, and STS-B datasets. For the SST-2, MNLI,
QNLI, and RTE datasets, we report the validation accuracy. For CoLA, we use
Matthew’s correlation as the evaluation metric. The F1 score is reported for
both MRPC and QQP tasks, while the Pearson-Spearman correlation is used to
evaluate the performance on the STS-B dataset. To evaluate the memory usage,
we report the peak GPU memory usage and compression rate during the fine-
tuning process with Huggingface API (Wolf et al., 2020).

Compared Methods and Adopted Models. We consider three methods to compare in
this paper: Full fine-tuning (Full), LoRA (Hu et al., 2021), and Ladder Side-
tuning (LST) (Sung et al., 2022). Specifically, Full tunes all of the
parameters in the model to provide an 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, we 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, we follow the previous work (Sung et al., 2022; Houlsby et al., 2019;
Hu et al., 2021) to adopt the Bert-Base (Devlin et al., 2018), Bert-Large,
T5-Base, T5-Large, and T5-3B (Raffel et al., 2020a) for evaluating the
effectiveness of different methods.

Hyperparameter Settings. For WTA-CRS , it only has one hyperparameter $k$,
which controls the column-row pair budget. We assign the same $k$ to all
replaceable linear operations in the model. We consider the normalized column-
row pair budget $k/|\mathcal{D}|\in\\{0.3,0.1\\}$, which are denoted as WTA-
CRS<EMAIL_ADDRESS>and WTA-CRS<EMAIL_ADDRESS>respectively. We also consider the combination
of 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.

Table 1: The GLUE benchmark results with T5 and Bert at different scales.

Model | Method | CoLA | SST-2 | MRPC | QQP | MNLI | QNLI | RTE | STS-B | AVG
---|---|---|---|---|---|---|---|---|---|---
BERT-Base | Full | 60.9$\pm$1.89 | 92.2$\pm$0.34 | 87.9$\pm$0.46 | 87.8$\pm$0.01 | 83.7$\pm$0.05 | 90.7$\pm$0.14 | 66.4$\pm$0.36 | 88.1$\pm$0.27 | 82.2
LoRA | 61.6$\pm$0.25 | 91.7$\pm$0.17 | 90.0$\pm$0.34 | 86.9$\pm$0.1 | 83.6$\pm$0.02 | 90.8$\pm$0.17 | 68.2$\pm$0.36 | 87.6$\pm$0.52 | 82.6
<EMAIL_ADDRESS>| 60.7$\pm$0.89 | 90.2$\pm$0.06 | 87.0$\pm$0.14 | 87.5$\pm$0.03 | 83.4$\pm$0.06 | 90.4$\pm$0.07 | 65.9$\pm$0.18 | 89.3$\pm$0.2 | 81.8
<EMAIL_ADDRESS>| 61.5$\pm$1.08 | 89.6$\pm$0.52 | 89.6$\pm$0.09 | 86.3$\pm$0.02 | 82.8$\pm$0.35 | 90.6$\pm$0.16 | 67.9$\pm$0.72 | 87.3$\pm$0.7 | 81.9
T5-Base | Full | 60.1$\pm$0.37 | 94.9$\pm$0.29 | 91.5$\pm$0.29 | 88.5$\pm$0.07 | 87.0$\pm$0.1 | 93.3$\pm$0.03 | 79.4$\pm$0.78 | 90.6$\pm$0.14 | 85.7
LoRA | 60.6$\pm$0.94 | 60.6$\pm$0.94 | 92.2$\pm$0.31 | 87.4$\pm$0.06 | 86.2$\pm$0.06 | 93.4$\pm$0.03 | 80.6$\pm$0.74 | 90.7$\pm$0.05 | 85.7
LST | 55.5$\pm$0.24 | 94.0$\pm$0.17 | 91.1$\pm$0.18 | 87.4$\pm$0.03 | 85.7$\pm$0.13 | 93.4$\pm$0.0 | 72.7$\pm$0.54 | 90.4$\pm$0.06 | 83.8
<EMAIL_ADDRESS>| 60.9$\pm$0.52 | 94.8$\pm$0.14 | 91.1$\pm$0.35 | 88.0$\pm$0.11 | 86.3$\pm$0.02 | 93.1$\pm$0.07 | 78.7$\pm$0.59 | 90.5$\pm$0.05 | 85.4
<EMAIL_ADDRESS>| 60.0$\pm$0.51 | 94.4$\pm$0.16 | 92.0$\pm$0.38 | 87.3$\pm$0.04 | 85.6$\pm$0.08 | 93.2$\pm$0.01 | 80.1$\pm$1.02 | 90.4$\pm$0.06 | 85.4
BERT-Large | Full | 66.8$\pm$0.31 | 93.5$\pm$0.29 | 89.5$\pm$0.26 | 88.5$\pm$0.03 | 86.4$\pm$0.19 | 92.1$\pm$0.24 | 72.6$\pm$0.36 | 90.2$\pm$0.76 | 85.0
LoRA | 65.9$\pm$0.27 | 93.8$\pm$0.17 | 90.8$\pm$0.37 | 87.6$\pm$0.08 | 85.9$\pm$0.05 | 92.0$\pm$0.2 | 71.3$\pm$0.18 | 90.3$\pm$0.09 | 84.7
<EMAIL_ADDRESS>| 64.7$\pm$0.44 | 93.5$\pm$0.0 | 89.3$\pm$0.39 | 88.2$\pm$0.04 | 85.2$\pm$0.03 | 91.9$\pm$0.12 | 73.8$\pm$0.54 | 90.4$\pm$0.02 | 84.6
<EMAIL_ADDRESS>| 66.0$\pm$0.33 | 93.3$\pm$0.29 | 89.7$\pm$1.32 | 87.6$\pm$0.02 | 86.0$\pm$0.07 | 91.9$\pm$0.14 | 72.4$\pm$0.17 | 89.7$\pm$0.04 | 84.6
T5-Large | Full | 61.3$\pm$1.01 | 96.3$\pm$0.0 | 93.4$\pm$0.13 | 89.7$\pm$0.01 | 89.8$\pm$0.07 | 94.2$\pm$0.05 | 85.3$\pm$0.17 | 91.8$\pm$0.08 | 87.7
LoRA | 63.3$\pm$0.26 | 96.4$\pm$0.14 | 93.5$\pm$0.16 | 88.5$\pm$0.03 | 89.5$\pm$0.05 | 94.3$\pm$0.07 | 84.2$\pm$0.68 | 91.7$\pm$0.13 | 87.7
LST | 59.9$\pm$0.77 | 95.8$\pm$0.06 | 91.8$\pm$0.08 | 88.4$\pm$0.01 | 88.7$\pm$0.05 | 94.2$\pm$0.02 | 82.5$\pm$0.18 | 91.4$\pm$0.07 | 86.6
<EMAIL_ADDRESS>| 60.9$\pm$1.18 | 96.3$\pm$0.25 | 93.6$\pm$0.57 | 89.3$\pm$0.04 | 89.5$\pm$0.12 | 94.1$\pm$0.03 | 84.4$\pm$0.34 | 91.3$\pm$0.05 | 87.4
<EMAIL_ADDRESS>| 62.9$\pm$1.19 | 96.2$\pm$0.05 | 93.6$\pm$0.47 | 88.3$\pm$0.02 | 89.2$\pm$0.08 | 94.0$\pm$0.07 | 83.9$\pm$0.95 | 91.3$\pm$0.03 | 87.4
T5-3B | LoRA | 70.1$\pm$0.37 | 96.8$\pm$0.29 | 94.0$\pm$0.27 | 89.9$\pm$0.0 | 91.0$\pm$0.14 | 95.6$\pm$0.05 | 85.9$\pm$0.36 | 92.9$\pm$0.08 | 89.5
<EMAIL_ADDRESS>| 71.4$\pm$0.35 | 96.4$\pm$0.06 | 94.6$\pm$0.39 | 90.0$\pm$0.05 | 91.0$\pm$0.06 | 95.6$\pm$0.12 | 86.3$\pm$0.36 | 92.9$\pm$0.09 | 89.8

Table 2: Peak memory usage (GB) and compression rate of fine-tuning T5-Base
and -Large. We measure the memory usage on a 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 6 for it.

| FP | LoRA | LST |  WTA-CRS<EMAIL_ADDRESS>|  WTA-CRS<EMAIL_ADDRESS>|  LoRA+WTA-CRS<EMAIL_ADDRESS>|  LoRA+WTA-CRS<EMAIL_ADDRESS>
---|---|---|---|---|---|---|---
T5-Base |  17.66 (1$\times$) |  13.84 (1.3$\times$) |  5.50 (3.2$\times$) |  8.44 (2.1$\times$) |  7.30 (2.4$\times$) |  6.50 (2.7$\times$) |  5.44 (3.2$\times$)
T5-Large |  45.85 (1$\times$) |  36.83 (1.2$\times$) |  14.85 (3.1$\times$) |  21.58 (2.1$\times$) |  18.46 (2.5$\times$) |  17.44 (2.6$\times$) |  14.64 (3.13$\times$)

###  5.2 Accuracy versus Memory Usage (RQ1)

To answer RQ1, we first analyze the trade-off between the model performance
and memory saving. The evaluation results and peak memory usage are given in
Tables 1 and 2, respectively. We observe:

❶ _WTA-CRS achieves a superior trade-off between accuracy and memory usage
compared to baselines. Specifically, WTA-CRS has negligible accuracy drop,
while the peak memory usage is reduced by $2.1\times\sim 2.7\times$_ (when
combined with LoRA).

Figure 6: Peak memory usage versus the maximum batch size of T5-3B. More
similar results are shown in Appendix E.2.

As we analyzed, LoRA mainly reduces the memory of optimizer states. Thus,
although it has negligible accuracy drop, it can only achieve $\sim 1.3\times$
peak memory saving. LST can reduce memory usage up to $3\times$, but its
accuracy drop is much larger than LoRA and WTA-CRS . Since WTA-CRS executes at
the operation level and focuses on the activations, we further combine LoRA
with WTA-CRS to reduce memory usage more aggressively. When combing with LoRA,
WTA-CRS achieves $2.7\times$ memory usage saving with almost no accuracy drop.
To fully tune T5-3B, it requires 37.7GB of memory and relies on a GPU with a
capacity of 40GB or higher, e.g. RTX8000, A100, or A40. On the other hand,
LoRA+WTA-CRS only requires 21.6GB of memory for finetuning with a mini-batch
size of $32$, which can run on a 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 1, ❷ _under the
similar memory budget, WTA-CRS outperforms the other methods in terms of the
accuracy._

Also, according to Figure 6, for T5-3B, LoRA itself can enable $1.9\times$
larger batch size. ❸ _When combined with LoRA, WTA-CRS enables $4.8\times$
($k\\!=\\!0.3|\mathcal{D}|$) to $6.4\times$ ($k\\!=\\!0.1|\mathcal{D}|$)
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 ❹ _It has almost no accuracy
drop when $k=0.3|\mathcal{D}|$. And the accuracy drop is about $1\%$ when
$k=0.1|\mathcal{D}|$_. _Notably, for T5-3B, the accuracy drop is only $0.4\%$
when $k=0.1|\mathcal{D}|$, 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 (Li et
al., 2020).

(a) T5-Base

(b) T5-Large

(c) T5-3B

Figure 7: Average validation results on GLUE dataset of WTA-CRS with varying
budgets.

###  5.3 Ablation Study (RQ3 and RQ4)

(a) SST2

(b) MNLI

(c) QQP

Figure 8: Validation results of T5-Base with different methods.

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 $k$ probability of Equation (3). We note that this is
the estimator proposed in (Adelman et al., 2021). (2) The CRS method follows
Equation (3) to sample the row-column pairs. All methods are deployed to 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/|\mathcal{D}|=0.1$ for all
methods. The validation accuracy versus training epoch is given in Figure 8.
We observe:

❺ _WTA-CRS outperforms all compared methods, especially as the training epoch
grows._ The deterministic selection of top $k$ column-row pairs suffers from
accumulation of bias error that ultimately results in a failure of
convergence. For CRS, it 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 to larger variance. In summary, both the
deterministic and stochastic parts contribute to the effectiveness of WTA-CRS
, which is consistent with our theoretical analysis.

(a) T5-Large

(b) T5-3B

Figure 9: Batch size versus training throughput (sentences/sec) with different
methods, where the sequential length is 128. The hardware is one single
NVIDIA-A100 (80GB).

#### The speed of WTA-CRS (RQ4).

The configuration of computational infrastructure is given in Appendix F.1. 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.
The memory taken by the activations is proportional to the batch size and
sequential length. Below we analyze how the fine-tuning speed affected by WTA-
CRS . As we analyzed in Appendix A limitation, the current implementation 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, a reduction in activation memory enables the
use of larger batch sizes, thereby improving training speed due to increased
GPU utilization (Goyal et al., 2017; Sung et al., 2022). As we analyzed in
Figure 6, WTA-CRS can enlarge the available batch size by up to $4.8\times$
larger. This enhancement is expected to result in a acceleration of the
training speed. To illustrate this relationship, Figure 9 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<EMAIL_ADDRESS>shows $1.08\times$ higher training throughput; and on the T5-3B
model, WTA-CRS<EMAIL_ADDRESS>and WTA-CRS<EMAIL_ADDRESS>achieve $1.14\times$ and $1.21\times$
higher training throughput, respectively.

##  6 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.7\times$, and it enables a $6.4\times$ larger batch size, which in return
resulting in $1.2\times$ higher training throughput.

## References

  * Abadi et al. (2016) Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, et al.  Tensorflow: a system for large-scale machine learning.  In _Osdi_ , volume 16, pages 265–283. Savannah, GA, USA, 2016.
  * Adelman et al. (2021) Menachem Adelman, Kfir Levy, Ido Hakimi, and Mark Silberstein.  Faster neural network training with approximate tensor operations.  _Advances in Neural Information Processing Systems_ , 34:27877–27889, 2021.
  * Andoorveedu et al. (2022a) Muralidhar Andoorveedu, Zhanda Zhu, Bojian Zheng, and Gennady Pekhimenko.  Tempo: Accelerating transformer-based model training through memory footprint reduction.  _arXiv preprint arXiv:2210.10246_ , 2022a.
  * Andoorveedu et al. (2022b) Muralidhar Andoorveedu, Zhanda Zhu, Bojian Zheng, and Gennady Pekhimenko.  Tempo: Accelerating transformer-based model training through memory footprint reduction.  _arXiv preprint arXiv:2210.10246_ , 2022b.
  * Brown et al. (2020) 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–1901, 2020.
  * Chen et al. (2021a) 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:2112.00029_ , 2021a.
  * Chen et al. (2021b) Beidi Chen, Zichang Liu, Binghui Peng, Zhaozhuo Xu, Jonathan Lingjie Li, Tri Dao, Zhao Song, Anshumali Shrivastava, and Christopher Re.  MONGOOSE: A learnable LSH framework for efficient neural network training.  In _International Conference on Learning Representations (ICLR)_ , 2021b.  URL https://openreview.net/forum?id=wWK7yXkULyh.
  * Chen et al. (2021c) 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, 2021c.
  * Chmiel et al. (2021) 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.
  * Chmiel et al. (2023) Brian Chmiel, Itay Hubara, Ron Banner, and Daniel Soudry.  Minimum variance unbiased n: M sparsity for the neural gradients.  In _The Eleventh International Conference on Learning Representations_ , 2023.
  * Chuang et al. (2023) Yu-Neng Chuang, Guanchu Wang, Fan Yang, Zirui Liu, Xuanting Cai, Mengnan Du, and Xia Hu.  Efficient XAI techniques: A taxonomic survey.  _CoRR_ , abs/2302.03225, 2023.  doi: 10.48550/arXiv.2302.03225.  URL https://doi.org/10.48550/arXiv.2302.03225.
  * Dao et al. (2022) 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 on Machine Learning_ , pages 4690–4721. PMLR, 2022.
  * Devlin et al. (2018) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  Bert: Pre-training of deep bidirectional transformers for language understanding.  _arXiv preprint arXiv:1810.04805_ , 2018.
  * Drineas and Kannan (2001) Petros Drineas and Ravi Kannan.  Fast monte-carlo algorithms for approximate matrix multiplication.  In _Proceedings 42nd IEEE Symposium on Foundations of Computer Science_ , pages 452–459. IEEE, 2001.
  * Drineas et al. (2006) 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\.
  * Goyal et al. (2017) 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 1 hour.  _arXiv preprint arXiv:1706.02677_ , 2017.
  * Houlsby et al. (2019) Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin De Laroussilhe, Andrea Gesmundo, Mona Attariyan, and Sylvain Gelly.  Parameter-efficient transfer learning for nlp.  In _International Conference on Machine Learning_ , pages 2790–2799. PMLR, 2019.
  * Hu et al. (2021) Edward J 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.09685_ , 2021.
  * Jain et al. (2020) 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 rematerialization.  _Proceedings of Machine Learning and Systems_ , 2:497–511, 2020.
  * Kaplan et al. (2020) 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.
  * Karimi Mahabadi et al. (2021) Rabeeh Karimi Mahabadi, James Henderson, and Sebastian Ruder.  Compacter: Efficient low-rank hypercomplex adapter layers.  _Advances in Neural Information Processing Systems_ , 34:1022–1035, 2021.
  * Kingma and Ba (2014) Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ , 2014.
  * Kirisame et al. (2020) Marisa Kirisame, Steven Lyubomirsky, Altan Haan, Jennifer Brennan, Mike He, Jared Roesch, Tianqi Chen, and Zachary Tatlock.  Dynamic tensor rematerialization.  _arXiv preprint arXiv:2006.09616_ , 2020.
  * Lester et al. (2021) Brian Lester, Rami Al-Rfou, and Noah Constant.  The power of scale for parameter-efficient prompt tuning.  _arXiv preprint arXiv:2104.08691_ , 2021.
  * Li and Liang (2021) Xiang Lisa Li and Percy Liang.  Prefix-tuning: Optimizing continuous prompts for generation.  _arXiv preprint arXiv:2101.00190_ , 2021.
  * Li et al. (2020) Zhuohan Li, Eric Wallace, Sheng Shen, Kevin Lin, Kurt Keutzer, Dan Klein, and Joey Gonzalez.  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.
  * Liu et al. (2022a) 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, 2022a.
  * Liu et al. (2020) Zirui Liu, Qingquan Song, Kaixiong Zhou, Ting-Hsiang Wang, Ying Shan, and Xia Hu.  Towards interaction detection using topological analysis on neural networks.  _CoRR_ , abs/2010.13015, 2020.  URL https://arxiv.org/abs/2010.13015.
  * Liu et al. (2021) Zirui Liu, Haifeng Jin, Ting-Hsiang Wang, Kaixiong Zhou, and Xia Hu.  Divaug: plug-in automated data augmentation with explicit diversity maximization.  In _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , pages 4762–4770, 2021.
  * Liu et al. (2022b) Zirui Liu, Shengyuan Chen, Kaixiong Zhou, Daochen Zha, Xiao Huang, and Xia Hu.  Rsc: Accelerating graph neural networks training via randomized sparse computations.  _arXiv preprint arXiv:2210.10737_ , 2022b.
  * Liu et al. (2022c) 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 on Learning Representations_ , 2022c.  URL https://openreview.net/forum?id=vkaMaq95_rX.
  * Oktay et al. (2020) Deniz Oktay, Nick McGreivy, Joshua Aduol, Alex Beatson, and Ryan P Adams.  Randomized automatic differentiation.  _arXiv preprint arXiv:2007.10412_ , 2020.
  * Paszke et al. (2019) 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.
  * Raffel et al. (2020a) 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 a unified text-to-text transformer.  _The Journal of Machine Learning Research_ , 21(1):5485–5551, 2020a.
  * Raffel et al. (2020b) 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 a unified text-to-text transformer.  _The Journal of Machine Learning Research_ , 21(1):5485–5551, 2020b.
  * Sung et al. (2022) Yi-Lin Sung, Jaemin Cho, and Mohit Bansal.  Lst: Ladder side-tuning for parameter and memory efficient transfer learning.  _arXiv preprint arXiv:2206.06522_ , 2022.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  _Advances in neural information processing systems_ , 30, 2017.
  * Wang et al. (2018a) 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_ , 2018a.
  * Wang et al. (2018b) 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_ , 2018b.
  * Wang et al. (2022a) 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, 2022a.
  * Wang et al. (2022b) 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.04187_ , 2022b.
  * Wolf et al. (2020) Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, et al.  Transformers: State-of-the-art natural language processing.  In _Proceedings of the 2020 conference on empirical methods in natural language processing: system demonstrations_ , pages 38–45, 2020.
  * Yang et al. (2023) Jingfeng Yang, Hongye Jin, Ruixiang Tang, Xiaotian Han, Qizhang Feng, Haoming Jiang, Bing Yin, and Xia Hu.  Harnessing the power of llms in practice: A survey on chatgpt and beyond.  _arXiv preprint arXiv:2304.13712_ , 2023.
  * Zaken et al. (2021) Elad Ben Zaken, Shauli Ravfogel, and Yoav Goldberg.  Bitfit: Simple parameter-efficient fine-tuning for transformer-based masked language-models.  _arXiv preprint arXiv:2106.10199_ , 2021.
  * Zha et al. (2023) Daochen Zha, Zaid Pervaiz Bhat, Kwei-Herng Lai, Fan Yang, Zhimeng Jiang, Shaochen Zhong, and Xia Hu.  Data-centric artificial intelligence: A survey.  _arXiv preprint arXiv:2303.10158_ , 2023.
  * Zhong et al. (2022) Shaochen Zhong, Guanqun Zhang, Ningjia Huang, and Shuai Xu.  Revisit kernel pruning with lottery regulated grouped convolutions.  In _International Conference on Learning Representations_ , 2022.  URL https://openreview.net/forum?id=LdEhiMG9WLO.

## Appendix

##  Appendix A Extended Related Work and Discussion

#### Parameter-Efficient Fine-tuning.

Parameter-efficient tuning methods select a small subset of parameters or
insert a few parameters to a pre-trained network. Then they only update the
small subset of parameters, while keeping others fixed (Lester et al., 2021;
Sung et al., 2022; Li and Liang, 2021; Zaken et al., 2021; Hu et al., 2021;
Karimi Mahabadi et al., 2021; Houlsby et al., 2019; Chen et al., 2021b). For
example, Adapters (Houlsby et al., 2019; Karimi Mahabadi et al., 2021) insert
a small module into the transformer blocks and only update it. Similarly,
prompt tuning (Lester et al., 2021) introduces a small vector that is
concatenated with the input embeddings. BitFit (Zaken et al., 2021) only tunes
the bias term of the model. LoRA (Hu et al., 2021) injects trainable rank
decomposition matrices into the transformer block. Although these methods are
“parameter-efficient”, they actually cannot reduce the memory usage of the
model itself. This is because we still needs to build the computation graph
for the whole model. Instead, the memory usage of optimizer states will be
significantly reduced, which is in proportional to the number of trainable
parameters (Kingma and Ba, 2014).

#### 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 (Kirisame et
al., 2020; Jain et al., 2020). According to the report of
Checkmate222https://github.com/parasj/checkmate/issues/153 (Jain et al.,
2020), it achieves “a 2.3x memory reduction when training a BERT model with
Checkmate optimizations (at 1x 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-fusion
technique (Liu et al., 2022a); (2) the existing implementation can still yield
a large speedup when employing larger batch sizes (Figure 9).

#### 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.

##  Appendix B Unbiasedness of Weight Gradient

This part we directly follow the proof of Theorem 1 in ActNN (Chen et al.,
2021c). For completeness, we provide the proof sketch here that is short and
easy to follow. Specifically, here we use ReLU as 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 $l^{\text{th}}$ layer is

$\displaystyle{\mathbf{Z}}^{(l+1)}$
$\displaystyle={\mathbf{H}}^{(l)}{\mathbf{W}}^{(l)},$
$\displaystyle{\mathbf{H}}^{(l+1)}$
$\displaystyle=\textrm{ReLU}({\mathbf{Z}}^{(l+1)}),$

and the backward pass of ReLU is:

$\displaystyle\mathbb{E}[\nabla{\mathbf{Z}}^{(l+1)}]$
$\displaystyle=\mathbb{E}[\mathbbm{1}_{{\mathbf{Z}}^{(l+1)}>0}\odot\nabla{\mathbf{H}}^{(l+1)}]$
$\displaystyle=\mathbbm{1}_{{\mathbf{Z}}^{(l+1)}>0}\odot\mathbb{E}[\nabla{\mathbf{H}}^{(l+1)}],$

where $\odot$ is the element-wise product and $\mathbbm{1}$ is the indicator
function. The element-wise product is linear operation and
$\mathbbm{1}_{{\mathbf{Z}}^{(l+1)}>0}$ is only related to the pre-activation
${\mathbf{Z}}^{(l+1)}$ in the forward pass. We only apply the approximation
during the backward pass so $\mathbbm{1}_{{\mathbf{Z}}^{(l+1)}>0}$ can be
extracted from the expectation. We know that for the last layer $L$, we have
$\mathbb{E}[\nabla{\mathbf{H}}^{(L)}]={\mathbf{H}}^{(L)}$ since we do not
apply activation at the output layer. We then can prove by induction that
$\mathbb{E}[\nabla{\mathbf{H}}^{(l+1)}]={\mathbf{H}}^{(l+1)}$ and
$\mathbb{E}[\nabla{\mathbf{W}}^{(l)}]={\mathbf{W}}^{(l)}$ for any layer $l$.

##  Appendix C Proof

###  C.1 Derivation of Equation (3)

Let ${\mathbf{X}}\in\mathbb{R}^{n\times m}$,
${\mathbf{Y}}\in\mathbb{R}^{m\times q}$ be two matrices. The matrix
multiplication ${\mathbf{X}}{\mathbf{Y}}$ can be estimated as

$\displaystyle\texttt{GEMM}({\mathbf{X}},{\mathbf{Y}})$
$\displaystyle=\sum_{i=1}^{m}{\mathbf{X}}_{:,i}{\mathbf{Y}}_{i,:}\approx\sum_{t=1}^{k}\frac{1}{kp_{i_{t}}}{\mathbf{X}}_{:,i_{t}}{\mathbf{Y}}_{i_{t},:}={\mathbf{X}}^{\prime}{\mathbf{Y}}^{\prime},$

Equation (3) shows the approximation error
$\mathbb{E}[||{\mathbf{X}}{\mathbf{Y}}-{\mathbf{X}}^{\prime}{\mathbf{Y}}^{\prime}||_{F}]$
is minimized when the probabilities

$p_{i}=\frac{||{\mathbf{X}}_{:,i}||_{2}\
||{\mathbf{Y}}_{i,:}||_{2}}{\sum_{j=1}^{m}||{\mathbf{X}}_{:,j}||_{2}\
||{\mathbf{Y}}_{j,:}||_{2}}.$

###### Proof.

Let
$f(i)=\frac{{\mathbf{X}}_{:i}{\mathbf{Y}}_{i:}}{p_{i}}\in\mathbb{R}^{n\times
q}$. We note that $f(i)$ is an unbiased estimation of
${\mathbf{X}}{\mathbf{Y}}$. Namely,

$\displaystyle\mathbb{E}_{j\sim\mathcal{P}}[f(j)]=\sum_{i=1}^{m}p_{i}\frac{{\mathbf{X}}_{:,i}{\mathbf{Y}}_{i:}}{p_{i}}={\mathbf{X}}{\mathbf{Y}}.$

Then we have

${\mathbf{X}}^{\prime}{\mathbf{Y}}^{\prime}=\frac{1}{k}\sum_{t=1}^{k}f(i_{t}),$
(8)

where $i_{1},\cdots,i_{t}$ are the index of the sampled column-row pairs at
$t^{\mathrm{th}}$ random trials. For each $i_{t}$, its variance is

$\displaystyle\mathrm{Var}[f(i_{t})]$
$\displaystyle=\mathrm{Var}[\frac{{\mathbf{X}}_{:i_{t}}{\mathbf{Y}}_{i_{t}:}}{p_{i_{t}}}]$
$\displaystyle=\mathbb{E}[\frac{{\mathbf{X}}^{2}_{:i_{t}}{\mathbf{Y}}^{2}_{i_{t}:}}{p^{2}_{i_{t}}}]-\mathbb{E}^{2}[\frac{{\mathbf{X}}_{:i_{t}}{\mathbf{Y}}_{i_{t}:}}{p_{i_{t}}}]$
$\displaystyle=\mathbb{E}[\frac{{\mathbf{X}}^{2}_{:i_{t}}{\mathbf{Y}}^{2}_{i_{t}:}}{p^{2}_{i_{t}}}]-({\mathbf{X}}{\mathbf{Y}})^{2}.$
$\displaystyle=\sum_{t=1}^{m}\frac{{\mathbf{X}}^{2}_{:t}{\mathbf{Y}}^{2}_{t:}}{p_{t}}-({\mathbf{X}}{\mathbf{Y}})^{2}.$
(9)

where the first step follows from the fact that
$\mathrm{Var}[{\mathbf{x}}]=\mathbb{E}[{\mathbf{x}}^{2}]-\mathbb{E}^{2}[{\mathbf{x}}]$.

Then we have,

$\displaystyle\mathbb{E}[||{\mathbf{X}}{\mathbf{Y}}-{\mathbf{X}}^{\prime}{\mathbf{Y}}^{\prime}||_{F}]$
$\displaystyle=\sum_{i=1}^{n}\sum_{j=1}^{q}\mathbb{E}[({\mathbf{X}}{\mathbf{Y}}-{\mathbf{X}}^{\prime}{\mathbf{Y}}^{\prime})_{ij}^{2}]$
$\displaystyle=\sum_{i=1}^{n}\sum_{j=1}^{q}\mathrm{Var}[({\mathbf{X}}^{\prime}{\mathbf{Y}}^{\prime})_{ij}].$

By combining Equation (8) and Equation (9) into the above equation, we have

$\displaystyle\mathbb{E}[||{\mathbf{X}}{\mathbf{Y}}-{\mathbf{X}}^{\prime}{\mathbf{Y}}^{\prime}||_{F}]$
$\displaystyle=\frac{1}{k}\sum_{i=1}^{n}\sum_{j=1}^{q}\sum_{t=1}^{m}\frac{{\mathbf{X}}^{2}_{it}{\mathbf{Y}}^{2}_{tj}}{p_{t}}-\frac{1}{k}\|{\mathbf{X}}{\mathbf{Y}}\|_{F}^{2}.$
$\displaystyle=\frac{1}{k}\sum_{t=1}^{m}\frac{\|{\mathbf{X}}_{:,t}\|_{2}^{2}\|{\mathbf{Y}}_{t,:}\|_{2}^{2}}{p_{t}}-\frac{1}{k}\|{\mathbf{X}}{\mathbf{Y}}\|_{F}^{2}.$

To minimize
$\mathbb{E}[||{\mathbf{X}}{\mathbf{Y}}-{\mathbf{X}}^{\prime}{\mathbf{Y}}^{\prime}||_{F}]$,
the optimal probability distribution can be obtained via solving the following
optimization problem:

$\displaystyle\min_{p_{1},\cdots,p_{m}}$
$\displaystyle\sum_{t=1}^{m}\frac{\|{\mathbf{X}}_{:,t}\|_{2}^{2}\|{\mathbf{Y}}_{t,:}\|_{2}^{2}}{p_{t}},$
s.t. $\displaystyle\sum_{t=1}^{m}p_{t}=1.$

The solution to the above convex problem is the distribution defined in
Equation (3). Namely,

$p_{i}=\frac{||{\mathbf{X}}_{:,i}||_{2}\
||{\mathbf{Y}}_{i,:}||_{2}}{\sum_{j=1}^{m}||{\mathbf{X}}_{:,j}||_{2}\
||{\mathbf{Y}}_{j,:}||_{2}}.$

∎

###  C.2 Unbiasedness of Our Proposed Estimator

See 1

###### Proof.

$\displaystyle\mathbb{E}_{j\sim\mathcal{P}^{\mathcal{D}\backslash\mathcal{C}}}\Big{[}\sum_{c\in\mathcal{C}}f(c)p_{c}+(1-\sum_{c\in\mathcal{C}}p_{c})f(j)\Big{]}$
$\displaystyle=$
$\displaystyle\sum_{c\in\mathcal{C}}f(c)p_{c}+(1-\sum_{c\in\mathcal{C}}p_{c})\mathbb{E}_{j\sim\mathcal{P}^{\mathcal{D}\backslash\mathcal{C}}}[f(j)]$
$\displaystyle=$
$\displaystyle\sum_{c\in\mathcal{C}}f(c)p_{c}+(1-\sum_{c\in\mathcal{C}}p_{c})\sum_{j\in\mathcal{D}\backslash\mathcal{C}}\frac{p_{j}}{1-\sum_{c\in\mathcal{C}}p_{c}}f(j)$
$\displaystyle=$
$\displaystyle\sum_{c\in\mathcal{C}}f(c)p_{c}+\sum_{j\in\mathcal{D}\backslash\mathcal{C}}f(j)p_{j}$
$\displaystyle=$ $\displaystyle\mathbb{E}_{j\sim\mathcal{P}}[f(j)]$
$\displaystyle=$ $\displaystyle{\mathbf{X}}{\mathbf{Y}}$

∎

###  C.3 Variance of Our Proposed Estimator

See 2

###### Proof.

Recall that the original estimator for matrix production
${\mathbf{X}}{\mathbf{Y}}$ is defined as

$\displaystyle\mathbb{E}_{i\sim\mathcal{P}}[f(i)].$ (10)

and our proposed family of estimator is defined as:

$\displaystyle
h(j)=\mathbb{E}_{j\sim\mathcal{P}^{\mathcal{D}\backslash\mathcal{C}}}\Big{[}\sum_{c\in\mathcal{C}}f(c)p_{c}+(1-\sum_{c\in\mathcal{C}}p_{c})f(j)\Big{]}.$
(11)

We first define three independent random variables as belows:

$\displaystyle u$ $\displaystyle\sim\mathcal{P}^{\mathcal{C}},$ (12)
$\displaystyle j$
$\displaystyle\sim\mathcal{P}^{\mathcal{D}\backslash\mathcal{C}},$ (13)
$\displaystyle b$
$\displaystyle\sim\text{Bernoulli}(1-\sum_{c\in\mathcal{C}}p_{c}).$ (14)

According to the Law of total variance, we have

$\displaystyle\mathrm{Var}[f(i)]$
$\displaystyle=\mathbb{E}_{b}\Big{[}Var[f(i)|b]\Big{]}+\mathrm{Var}_{b}\Big{[}\mathbb{E}[f(i)|b]\Big{]}$
$\displaystyle\geq\mathbb{E}_{b}\Big{[}\mathrm{Var}[f(i)|b]\Big{]}$
$\displaystyle=\sum_{c\in\mathcal{C}}p_{c}\mathrm{Var}[f(i)|b=0]+(1-\sum_{c\in\mathcal{C}}p_{c})\mathrm{Var}[f(i)|b=1]$
$\displaystyle\geq(1-\sum_{c\in\mathcal{C}}p_{c})\mathrm{Var}[f(i)|i\in\mathcal{D}\backslash\mathcal{C}]$
(15)

where the first step follows from the fact that for any random variance
${\mathbf{x}},{\mathbf{y}}$, we have
$\mathrm{Var}[{\mathbf{y}}]=\mathbb{E}[\mathrm{Var}[{\mathbf{y}}|{\mathbf{x}}]]+\mathrm{Var}[\mathbb{E}[{\mathbf{y}}|{\mathbf{x}}]]$.
Also, by Equation (11), we have

$\mathrm{Var}[h(j)]=(1-\sum_{c\in\mathcal{C}}p_{c})^{2}\mathrm{Var}[f(j)|j\in\mathcal{D}\backslash\mathcal{C}].$
(16)

By combining the above two inequality, we have

$\mathrm{Var}[h(j)]\leq(1-\sum_{c\in\mathcal{C}}p_{c})\mathrm{Var}[f(i)].$
(17)

Equation (17) quantitatively shows the variance reduction of $h(j)$ over
$f(i)$. Then we compare our estimator $\hat{g}({\mathbf{X}},{\mathbf{Y}})$ and
$g({\mathbf{X}},{\mathbf{Y}})$ in terms of variance.

First, because
$g({\mathbf{X}},{\mathbf{Y}})=\frac{1}{k}\sum_{t=1}^{k}f(i_{t}),~{}~{}~{}i_{1},\cdots
i_{k}\overset{\text{i.i.d}}{\sim}\mathcal{P}$. Thus we have

$\mathrm{Var}[g({\mathbf{X}},{\mathbf{Y}})]=\frac{1}{k}\mathrm{Var}[f(i)].$
(18)

Similarly, we have

$\mathrm{Var}[\hat{g}({\mathbf{X}},{\mathbf{Y}})]=\frac{1}{k-|\mathcal{C}|}\mathrm{Var}[h(j)].$
(19)

By combining Equation (17) into the above two equations, we have

$\displaystyle\mathrm{Var}[\hat{g}({\mathbf{X}},{\mathbf{Y}})]$
$\displaystyle=\frac{1}{k-|\mathcal{C}|}\mathrm{Var}[h(j)]$ (20)
$\displaystyle\leq\frac{1-\sum_{c\in\mathcal{C}}p_{c}}{k-|\mathcal{C}|}\mathrm{Var}[f(i)]$
$\displaystyle\leq\frac{1-\sum_{c\in\mathcal{C}}p_{c}}{k-|\mathcal{C}|}k\mathrm{Var}[g({\mathbf{X}},{\mathbf{Y}})],$

where the first step follows from Equation (17). By setting
$\frac{1-\sum_{c\in\mathcal{C}}p_{c}}{k-|\mathcal{C}|}k\leq 1$, we arrive the
conclusion that when $\sum_{c\in\mathcal{C}}p_{c}>\frac{|\mathcal{C}|}{k}$, we
have
$\mathrm{Var}[\hat{g}({\mathbf{X}},{\mathbf{Y}})]\leq\mathrm{Var}[g({\mathbf{X}},{\mathbf{Y}})]$.

Further, $\frac{1-\sum_{c\in\mathcal{C}}p_{c}}{k-|\mathcal{C}|}k$ achieves the
minimal when
$|\mathcal{C}|=\min_{|\mathcal{C}|\in\\{0,\cdots,k\\}}\frac{1-\sum_{c\in\mathcal{C}}p_{c}}{k-|\mathcal{C}|}$.

∎

##  Appendix D Implementation Details

The pseudocode for approximated linear layer with WTA-CRS and standard line
layer is given in Algorithm 1 and Algorithm 3, respectively. The column-row
pair sampling procedure is given in Algorithm 2. For the ease of illustration,
we ignore the sequential length. As we mentioned in the main text, we only
replace the GEMM in the backward pass with WTA-CRS . According to Equation
(1c), we need the activation gradient $\nabla{\mathbf{Z}}$ to perform the
column-row pair sampling during the forward pass. Thus we initialize a cache
in CPU memory to store the gradient norm of activations from the last step.
When performing column-row pair selection, we need to swap the gradient norm
of activations 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 a significant time overhead.

Hyperparameter: The total budget of column-row pairs $k$.

procedure  _Init_ :

Initialize $\texttt{Cache}\in\mathbb{R}^{N}$ as an empty matrix in main memory
// $N$ is the total number of samples in the dataset. Cache is used for saving
the norm of output gradient $\nabla{\mathbf{Z}}$.

end procedure

procedure  _Forward Pass_ :

Input: activation ${\mathbf{H}}\in\mathbb{R}^{B\times D}$, weight
${\mathbf{W}}\in\mathbb{R}^{D\times D}$, indices of the current batch samples
$BI=\\{j_{1},\cdots,j_{B}\\}$.

$\texttt{ctx}\leftarrow{\\{\\}}$ // the context which saves tensors for
backward

${\mathbf{Z}}={\mathbf{H}}{\mathbf{W}}$

${\mathbf{H}}^{\prime},ind\leftarrow$Subsample(${\mathbf{H}}$, Cache[$BI$],
$k$)

// Cache[$BI$] is the cached gradient norm from the backward pass; $ind$ is
the set of involved column-row pair indices

$\texttt{ctx}\leftarrow\\{{\mathbf{H}}^{\prime},{\mathbf{W}},BI,ind\\}$

return _${\mathbf{Z}}$_

end procedure

procedure  _Backward pass_ :

Input: ctx from the forward pass, output gradient
$\nabla{\mathbf{Z}}\in\mathbb{R}^{B\times D}$

${\mathbf{H}}^{\prime},{\mathbf{W}},BI,ind\leftarrow\texttt{ctx}$

$\nabla{\mathbf{H}}=\nabla{\mathbf{Z}}{\mathbf{W}}^{\top}$

$\nabla{\mathbf{Z}}^{\prime}\leftarrow\nabla{\mathbf{Z}}[ind]$

// $\nabla{\mathbf{Z}}^{\prime}\in\mathbb{R}^{k\times D}$

$\nabla{\mathbf{W}}={\mathbf{H}}^{\prime\top}\nabla{\mathbf{Z}}^{\prime}$

for  _j in $BI$_ do

Cache[$j$] = $\|\nabla{\mathbf{Z}}_{j,:}\|_{2}$

end for

// Update the gradient norm of samples in the current batch

return _$\nabla{\mathbf{H}},\nabla{\mathbf{W}}$_

end procedure

Algorithm 1 Forward & Backward pass of Approximated Linear Layer

Input: activation ${\mathbf{H}}\in\mathbb{R}^{B\times D}$, gradient norm
${\mathbf{z}}\in\mathbb{R}^{B}$, the total budget of column-row pairs $k$.

for  _$i=1,\cdots,B$_ do

$p_{i}\leftarrow\frac{{\mathbf{z}}_{i}||{\mathbf{H}}_{i,:}||_{2}}{\sum_{j=1}^{B}{\mathbf{z}}_{i}||{\mathbf{H}}_{j,:}||_{2}}$
// The probability of column-row pairs defined in Equation (3).

end for

$\hat{k}\leftarrow\min_{\hat{k}\in\\{0,\cdots,k\\}}\frac{1-\sum_{c\in\mathcal{C}}p_{c}}{k-\hat{k}}$,
s.t. $\mathcal{C}=|\hat{k}|$. // $\mathcal{C}$ is the set of column-row pair
indices associated with $|\mathcal{C}|$ largest $p_{i}$.

Sample $k-|\mathcal{C}|$ i.i.d. column-row pairs
$\mathcal{C}_{\text{stoc}}=\\{i_{1},\cdots,i_{k-|\mathcal{C}|}\\}$ from the
distribution $\mathcal{P}^{\mathcal{D}\backslash\mathcal{C}}$

$ind\leftarrow\mathcal{C}\cup\mathcal{C}_{\text{stoc}}$

for  _$j\in\mathcal{C}_{\text{stoc}}$_ do

${\mathbf{H}}[j,:]\leftarrow{\mathbf{H}}[j,:]*\frac{1-\sum_{c\in\mathcal{C}}p_{c}}{(k-|\mathcal{C}|)p_{j}}$
// We need to normalize the stochastic part in Equation (6) to ensure the
unbiasedness.

end for

${\mathbf{H}}^{\prime}\leftarrow{\mathbf{H}}[ind]$ //
${\mathbf{H}}^{\prime}\in\mathbb{R}^{k\times D}$

return _${\mathbf{H}}^{\prime}$ , ind_

Algorithm 2 Subsample

procedure  _Forward Pass_ :

Input: activation ${\mathbf{H}}_{Q}\in\mathbb{R}^{BS\times D}$, weight
${\mathbf{W}}_{Q}\in\mathbb{R}^{D\times D}$, batch indices $index$

$\texttt{ctx}\leftarrow{\\{\\}}$ // the context which saves tensors for
backward

${\mathbf{Z}}_{Q}={\mathbf{H}}_{Q}{\mathbf{W}}_{Q}$

$\texttt{ctx}\leftarrow\\{{\mathbf{H}}_{Q},{\mathbf{W}}_{Q}\\}$

return _${\mathbf{Z}}_{Q}$_

end procedure

procedure  _Backward pass_ :

Input: ctx from the forward pass, output gradient $\nabla{\mathbf{Z}}_{Q}$

${\mathbf{H}}_{Q},{\mathbf{W}}_{Q}\leftarrow\texttt{ctx}$

$\nabla{\mathbf{H}}_{Q}=\nabla{\mathbf{Z}}_{Q}{\mathbf{W}}_{Q}^{\top}$

$\nabla{\mathbf{W}}_{Q}={\mathbf{H}}_{Q}^{\top}\nabla{\mathbf{Z}}_{Q}$

return _$\nabla{\mathbf{H}}_{Q},\nabla{\mathbf{W}}_{Q}$_

end procedure

Algorithm 3 Forward & Backward pass of the standard Linear layer

##  Appendix E More Experimental Results

Figure 10: The probability mass $\sum_{c\in\mathcal{C}}p_{c}$ versus
$\frac{|\mathcal{C}|}{k}$ in Equation (7) at $k=0.1|\mathcal{D}|$. Here we
visualize the column-row index distribution of query/key/value layer T5-base
model, fine-tuned on RTE dataset.

Figure 11: The probability mass $\sum_{c\in\mathcal{C}}p_{c}$ versus
$\frac{|\mathcal{C}|}{k}$ in Equation (7) at $k=0.5|\mathcal{D}|$. Here we
visualize the column-row index distribution of query/key/value layer T5-base
model, fine-tuned on RTE dataset.

Figure 12: The probability mass of top-$10\%$ column-row pairs in Equation (3)
versus iterations. Here we visualize the query/key/value layer T5-base model,
fine-tuned on RTE dataset.

###  E.1 More Experimental Analysis on Theorem 2

To evaluate Theorem 2 more comprehensively, below we also plot the
$\sum_{c\in\mathcal{C}}p_{c}$ versus $\frac{|\mathcal{C}|}{k}$ at
$k=0.1|\mathcal{D}|$ and $k=0.5|\mathcal{D}|$ in Figure 10 and 11,
respectively. We also plot $\sum_{c\in\mathcal{C}}p_{c}$ versus iterations in
Figure 12. We summarize that the the column-row index distribution is
concentrated on a few column-row pairs. Thus, the assumption in Theorem 2
holds under the context of fine-tuning transformers.

###  E.2 More Experimental Speed Analysis

Increasing the batch size can often result in faster model convergence and/or
enhance the final performance. Ideally, we should adjust the batch size
according to the requirements of our model 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 13. 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 3, “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 1). However, we note
that (1) the overhead can be greatly reduced with better implementation, e.g.,
using prefetch and operation-fusion technique (Liu et al., 2022a); (2) the
existing implementation can still yield a large speedup when employing larger
batch sizes (Figure 9).

(a) T5-Base

(b) T5-Large

(c) T5-3B

Figure 13: Peak memory usage versus maximum mini-batch size of T5. | Method |  T5- ATT |  T5- FF |  T5 Block |  T5- Large
---|---|---|---|---|---
Fwd | Full | 8 | 10 | 17 | 1052
WTA-CRS  | 22 | 16 | 37 | 2013
Bwd | Full | 16 | 19 | 34 | 2073
WTA-CRS  | 15 | 14 | 30 | 1738
F-B | Full | 24 | 29 | 51 | 3125
WTA-CRS  | 37 | 30 | 67 | 3751
Table 3: Latency (ms) of Forward and Backward pass.

##  Appendix F Experimental Settings

We give the detailed hyper-parameter setting in this section. Specifically,
for both T5 and BERT models, the parameters are updated with the AdamW
optimizer $\beta_{1}=0.9$ $\beta_{2}=0.999$ $\epsilon=10^{-8}$ and
$\text{weight decay}=0$. The the learning rate is adjusted with a linear LR
Scheduler, which maintains a 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 a LoRA dimension 32 if 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 are given in Tables 5, 6, 7,
respectively.

###  F.1 Computational Infrastructure

The computational infrastructure information is given in Table 4.

Table 4: Computing infrastructure for the experiments. Device Attribute | Value
---|---
Computing infrastructure | GPU
GPU model | NVIDIA-A100
GPU Memory | 81251MB
CUDA Version | 11.4
CPU Memory | 512GB
Table 5: Learning rate.

Model | Method | CoLA | SST-2 | MRPC | QQP | MNLI | QNLI | RTE | STS-B
---|---|---|---|---|---|---|---|---|---
BERT-Base |  WTA-CRS<EMAIL_ADDRESS>| 2e-5
LoRA+WTA-CRS<EMAIL_ADDRESS>| 2e-4 | 5e-4 | 2e-4 | 3e-4 | 3e-4 | 2e-4 | 2e-4 | 3e-4
T5-Base |  WTA-CRS<EMAIL_ADDRESS>| 3e-5 | 3e-6 | 3e-5 | 3e-5
WTA-CRS<EMAIL_ADDRESS>| 3e-5
LoRA+WTA-CRS<EMAIL_ADDRESS>| 3e-4 | 3e-5 | 3e-4 | 3e-5 | 3e-5 | 3e-5 | 3e-4 | 3e-4
LoRA+WTA-CRS<EMAIL_ADDRESS>| 3e-4 | 3e-5 | 3e-4 | 3e-5 | 3e-5 | 3e-5 | 3e-4 | 3e-4
BERT-Large |  WTA-CRS<EMAIL_ADDRESS>| 2e-5
LoRA+WTA-CRS<EMAIL_ADDRESS>| 3e-4 | 2e-4 | 2e-4 | 2e-4 | 2e-4 | 2e-4 | 3e-4 | 3e-4
T5-Large |  WTA-CRS<EMAIL_ADDRESS>| 3e-5 | 3e-6 | 3e-5 | 3e-5
WTA-CRS<EMAIL_ADDRESS>| 3e-5 | 3e-6 | 3e-5 | 3e-5
| LoRA+WTA-CRS<EMAIL_ADDRESS>| 3e-4 | 3e-5 | 3e-4 | 3e-5 | 3e-5 | 3e-5 | 3e-4 | 3e-4
| LoRA+WTA-CRS<EMAIL_ADDRESS>| 3e-4 | 3e-5 | 3e-4 | 3e-5 | 3e-5 | 3e-5 | 3e-4 | 3e-4
T5-3B | LoRA+WTA-CRS<EMAIL_ADDRESS>| 3e-4 | 3e-5 | 3e-4 | 3e-4 | 3e-4 | 3e-5 | 3e-4 | 3e-4
LoRA+WTA-CRS<EMAIL_ADDRESS>| 3e-4 | 3e-5 | 3e-4 | 3e-4 | 3e-4 | 3e-5 | 3e-4 | 3e-4

Table 6: Training epoch number.

Model | Method | CoLA | SST-2 | MRPC | QQP | MNLI | QNLI | RTE | STS-B
---|---|---|---|---|---|---|---|---|---
BERT-Base |  WTA-CRS<EMAIL_ADDRESS>| 20 | 20 | 10 | 10 | 10 | 10 | 20 | 10
LoRA+WTA-CRS<EMAIL_ADDRESS>| 60 | 20 | 20 | 20 | 20 | 20 | 40 | 40
T5-Base |  WTA-CRS<EMAIL_ADDRESS>| 40 | 10 | 20 | 10 | 10 | 10 | 50 | 20
WTA-CRS<EMAIL_ADDRESS>| 40 | 10 | 20 | 10 | 10 | 10 | 50 | 20
LoRA+WTA-CRS<EMAIL_ADDRESS>| 40 | 10 | 20 | 20 | 20 | 10 | 50 | 20
LoRA+WTA-CRS<EMAIL_ADDRESS>| 40 | 10 | 20 | 20 | 20 | 10 | 50 | 20
BERT-Large |  WTA-CRS<EMAIL_ADDRESS>| 60 | 20 | 20 | 10 | 10 | 10 | 40 | 10
LoRA+WTA-CRS<EMAIL_ADDRESS>| 60 | 20 | 20 | 20 | 20 | 20 | 40 | 40
T5-Large |  WTA-CRS<EMAIL_ADDRESS>| 20 | 10 | 20 | 10 | 10 | 10 | 40 | 20
WTA-CRS<EMAIL_ADDRESS>| 20 | 10 | 20 | 10 | 10 | 10 | 40 | 20
LoRA+WTA-CRS<EMAIL_ADDRESS>| 40 | 10 | 40 | 10 | 10 | 10 | 60 | 20
LoRA+WTA-CRS<EMAIL_ADDRESS>| 40 | 10 | 20 | 10 | 10 | 10 | 60 | 20
T5-3B | LoRA+WTA-CRS<EMAIL_ADDRESS>| 40 | 10 | 20 | 10 | 10 | 10 | 60 | 20
LoRA+WTA-CRS<EMAIL_ADDRESS>| 40 | 10 | 20 | 10 | 10 | 10 | 60 | 20

Table 7: Training mini-batch size.

Model | Method | CoLA | SST-2 | MRPC | QQP | MNLI | QNLI | RTE | STS-B
---|---|---|---|---|---|---|---|---|---
BERT-Base/Large |  WTA-CRS<EMAIL_ADDRESS>| 128 | 16
LoRA+WTA-CRS<EMAIL_ADDRESS>| 128 | 16
T5-Base/Large/3B |  WTA-CRS<EMAIL_ADDRESS>| 100
WTA-CRS<EMAIL_ADDRESS>| 100
LoRA+WTA-CRS<EMAIL_ADDRESS>| 100
LoRA+WTA-CRS<EMAIL_ADDRESS>| 100