SlowGuess/NeurIPS_2023_OCR huggingface dataset:
To minimize  $\mathbb{E}[\| \mathbf{X}\mathbf{Y} - \mathbf{X}'\mathbf{Y}'\| _F]$ , the optimal probability distribution can be obtained via solving the following optimization problem:
$$
\min  _ {p _ {1}, \dots , p _ {m}} \sum_ {t = 1} ^ {m} \frac {\| \mathbf {X} _ {: , t} \| _ {2} ^ {2} \| \mathbf {Y} _ {t , :} \| _ {2} ^ {2}}{p _ {t}},
$$
$$
\begin{array}{l} \text {s . t .} \sum_ {t = 1} ^ {m} p _ {t} = 1. \end{array}
$$
The solution to the above convex problem is the distribution defined in Equation (3). Namely,
$$
p _ {i} = \frac {\left| \left| \mathbf {X} _ {: , i} \right| \right| _ {2} \left| \left| \mathbf {Y} _ {i , :} \right| \right| _ {2}}{\sum_ {j = 1} ^ {m} \left| \left| \mathbf {X} _ {: , j} \right| \right| _ {2} \left| \left| \mathbf {Y} _ {j , :} \right| \right| _ {2}}.
$$
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  $\mathbf{X}\mathbf{Y}$ , i.e,  $\mathbb{E}_{j\sim \mathcal{P}^{\mathcal{D}}\backslash c}[\sum_{c\in \mathcal{C}}f(c)p_c + (1 - \sum_{c\in \mathcal{C}}p_c)f(j)] = \mathbf{X}\mathbf{Y}$ .
$$
\begin{array}{l} \mathbb {E} _ {j \sim \mathcal {P D} \backslash c} \left[ \sum_ {c \in \mathcal {C}} f (c) p _ {c} + \left(1 - \sum_ {c \in \mathcal {C}} p _ {c}\right) f (j) \right] \\ = \sum_ {c \in \mathcal {C}} f (c) p _ {c} + (1 - \sum_ {c \in \mathcal {C}} p _ {c}) \mathbb {E} _ {j \sim \mathcal {P D} \backslash c} [ f (j) ] \\ = \sum_ {c \in \mathcal {C}} f (c) p _ {c} + \left(1 - \sum_ {c \in \mathcal {C}} p _ {c}\right) \sum_ {j \in \mathcal {D} \backslash \mathcal {C}} \frac {p _ {j}}{1 - \sum_ {c \in \mathcal {C}} p _ {c}} f (j) \\ = \sum_ {c \in \mathcal {C}} f (c) p _ {c} + \sum_ {j \in \mathcal {D} \backslash \mathcal {C}} f (j) p _ {j} \\ = \mathbb {E} _ {j \sim \mathcal {P}} [ f (j) ] \\ = X Y \\ \end{array}
$$
C.3 Variance of Our Proposed Estimator
Proof. Recall that the original estimator for matrix production  $\mathbf{X}\mathbf{Y}$  is defined as
$$
\mathbb {E} _ {i \sim \mathcal {P}} [ f (i) ]. \tag {10}
$$
and our proposed family of estimator is defined as:
$$
h (j) = \mathbb {E} _ {j \sim \mathcal {P} ^ {\mathcal {D}} \backslash c} \left[ \sum_ {c \in \mathcal {C}} f (c) p _ {c} + \left(1 - \sum_ {c \in \mathcal {C}} p _ {c}\right) f (j) \right]. \tag {11}
$$
We first define three independent random variables as belows:
$$
u \sim \mathcal {P} ^ {\mathcal {C}}, \tag {12}
$$
$$
j \sim \mathcal {P} ^ {\mathcal {D} \backslash \mathcal {C}}, \tag {13}
$$
$$
b \sim \operatorname {B e r n o u l l i} \left(1 - \sum_ {c \in \mathcal {C}} p _ {c}\right). \tag {14}
$$
According to the Law of total variance, we have
$$
\begin{array}{l} \operatorname {V a r} [ f (i) ] = \mathbb {E} _ {b} \left[ \operatorname {V a r} [ f (i) | b ] \right] + \operatorname {V a r} _ {b} \left[ \mathbb {E} [ f (i) | b ] \right] \\ \geq \mathbb {E} _ {b} \left[ \operatorname {V a r} [ f (i) | b ] \right] \\ = \sum_ {c \in \mathcal {C}} p _ {c} \operatorname {V a r} [ f (i) | b = 0 ] + (1 - \sum_ {c \in \mathcal {C}} p _ {c}) \operatorname {V a r} [ f (i) | b = 1 ] \\ \geq \left(1 - \sum_ {c \in \mathcal {C}} p _ {c}\right) \operatorname {V a r} [ f (i) | i \in \mathcal {D} \backslash \mathcal {C} ] \tag {15} \\ \end{array}
$$
where the first step follows from the fact that for any random variance  $\mathbf{x},\mathbf{y}$ , we have  $\operatorname{Var}[\mathbf{y}] = \mathbb{E}[\operatorname{Var}[\mathbf{y}|\mathbf{x}]] + \operatorname{Var}[\mathbb{E}[\mathbf{y}|\mathbf{x}]]$ . Also, by Equation (11), we have
$$
\operatorname {V a r} [ h (j) ] = \left(1 - \sum_ {c \in \mathcal {C}} p _ {c}\right) ^ {2} \operatorname {V a r} [ f (j) | j \in \mathcal {D} \backslash \mathcal {C} ]. \tag {16}
$$
By combining the above two inequalities, we have
$$
\operatorname {V a r} [ h (j) ] \leq (1 - \sum_ {c \in \mathcal {C}} p _ {c}) \operatorname {V a r} [ f (i) ]. \tag {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, \dots, i_k \stackrel{\text{i.i.d.}}{\sim} \mathcal{P}$ . Thus we have
$$
\operatorname {V a r} [ g (\mathbf {X}, \mathbf {Y}) ] = \frac {1}{k} \operatorname {V a r} [ f (i) ]. \tag {18}
$$
$$
\operatorname {V a r} [ \hat {g} (\mathbf {X}, \mathbf {Y}) ] = \frac {1}{k - | \mathcal {C} |} \operatorname {V a r} [ h (j) ]. \tag {19}
$$
By combining Equation (17) into the above two equations, we have
$$
\begin{array}{l} \operatorname {V a r} [ \hat {g} (\mathbf {X}, \mathbf {Y}) ] = \frac {1}{k - | \mathcal {C} |} \operatorname {V a r} [ h (j) ] \tag {20} \\ \leq \frac {1 - \sum_ {c \in \mathcal {C}} p _ {c}}{k - | \mathcal {C} |} \operatorname {V a r} [ f (i) ] \\ \leq \frac {1 - \sum_ {c \in \mathcal {C}} p _ {c}}{k - | \mathcal {C} |} k \operatorname {V a r} [ g (\mathbf {X}, \mathbf {Y}) ], \\ \end{array}
$$
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 at the conclusion that when  $\sum_{c\in \mathcal{C}}p_c > \frac{|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,\dots ,k\}}\frac{1 - \sum_{c\in\mathcal{C}}p_c}{k - |\mathcal{C}|}$ .
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.
Algorithm 1: Forward & Backward pass of Approximated Linear Layer
Hyperparameter: The total budget of column-row pairs  $k$    
procedure INIT: Initialize 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 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,\dots ,j_B\}$  ctx  $\leftarrow \{\}$  // the context which saves tensors for backward  $\mathbf{Z} = \mathbf{H}\mathbf{W}$  H',ind  $\leftarrow$  SUBSAMPLE(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 ctx  $\leftarrow \{\mathbf{H}',\mathbf{W},\mathbf{B}\mathbf{I},\mathrm{ind}\}$  return Z   
end   
procedure BACKWARD PASS: Input: ctx from the forward pass, output gradient  $\nabla \mathbf{Z}\in \mathbb{R}^{B\times D}$  H', W, BI, ind  $\leftarrow$  ctx  $\nabla \mathbf{H} = \nabla \mathbf{Z}\mathbf{W}^{\top}$ $\nabla \mathbf{Z}'\gets \nabla \mathbf{Z}[ind]$  //  $\nabla \mathbf{Z}'\in \mathbb{R}^{k\times D}$ $\nabla \mathbf{W} = \mathbf{H}'\nabla \mathbf{Z}'$  for j in BI do | Cache[j]  $= \| \nabla Z_{j,:}\| _2$  end // Update the gradient norm of samples in the current batch return VH, VW
Algorithm 2: SUBSAMPLE
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,\dots ,B$  do   
 $\begin{array}{r}p_i\leftarrow \frac{\mathbf{z}_i||\mathbf{H}_{i,:}||_2}{\sum_{j = 1}^B\mathbf{z}_i||\mathbf{H}_{j,:}||_2} / / \end{array}$  The probability of column-row pairs defined in Equation (3).   
end   
 $\hat{k}\gets \min_{\hat{k}\in \{0,\dots ,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}_{\mathrm{stoc}} = \{i_1,\dots ,i_{k - |\mathcal{C}|}\}$  from the distribution  $\mathcal{P}^{\mathcal{D}}\backslash \mathcal{C}$  ind  $\leftarrow \mathcal{C}\cup \mathcal{C}_{\mathrm{stoc}}$    
for  $j\in \mathcal{C}_{\mathrm{stoc}}$  do   
 $\begin{array}{rl}{\mathbf{H}[j,:]\leftarrow \mathbf{H}[j,:]*\frac{1 - \sum_{c\in\mathcal{C}}p_c}{(k - |\mathcal{C}|)p_j}} & {/ / }\end{array}$  We need to normalize the stochastic part in Equation (6) to ensure the unbiasedness.   
end   
 $\mathbf{H}'\gets \mathbf{H}[ind]\quad /\quad \mathbf{H}'\in \mathbb{R}^{k\times D}$    
return  $\mathbf{H}'$  , ind
Algorithm 3: Forward & Backward pass of the standard Linear layer
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  
ctx  $\leftarrow \{\}$  // the context which saves tensors for backward  
 $\mathbf{Z}_Q = \mathbf{H}_Q \mathbf{W}_Q$   
ctx  $\leftarrow \{\mathbf{H}_Q, \mathbf{W}_Q\}$   
return  $\mathbf{Z}_Q$   
end  
procedure BACKWARD PASS:  
Input: ctx from the forward pass, output gradient  $\nabla \mathbf{Z}_Q$ $\mathbf{H}_Q, \mathbf{W}_Q \leftarrow \mathrm{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
Fig. 9. 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.
Fig. 10. The probability mass  $\sum_{c\in \mathcal{C}}p_c$  versus  $\frac{|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.
Fig. 11. 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.
(b) MNLI
(c) QQP
Fig. 12. Validation results of T5-Base with different methods.
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{|C|}{k}$  at  $k = 0.1|\mathcal{D}|$  and  $k = 0.5|\mathcal{D}|$  in Figure 9 and 10, respectively. We also plot  $\sum_{c\in \mathcal{C}}p_c$  versus iterations in Figure 11. We summarize that 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
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
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 [30]; (2) the existing implementation can still yield a large speedup when employing larger batch sizes (Figure 8).
(b) T5-Large
(c) T5-3B
Fig. 13. Peak memory usage versus maximum mini-batch size of T5.
E.3 Performance of WTA-CRS on OPT models
We conduct experiments of WTA-CRS on OPT models. It is observed in Table 4 that WTA-CRS exhibits 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.
Table 3: Latency (ms) of Forward and Backward pass.
Method	CoLA	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@0.3	51.8	88.43	74.01	86.61	75.2125
Table 4: Deployment of WTA-CRS@0.3 to OPT-350M.
E.4 Performance of Activation Quantization+WTA-CRS
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, on the T5-
base model. The experiment results are presented in Table 5. It is notable 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.
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-CRS@0.3+Quant@8bit	60.3	92.06	81.2	90.4	81.0
E.5 Comparison of 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 LoRA+WTA-CRS@0.3 over the quantized model. Below in Table 6, we demonstrate that applying QLoRA and WTA-CRS over a 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.
Table 5: Combination of wTA-CRS with quanzation on the T5-base model.
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 (NF4)+WTACRS@0.3	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.12
LoRA+WTA-CRS@0.3	6.50	17.44
QLoRA+WTA-CRS@0.3	6.31	16.75
Table 7: Peak memory usage (GB) of fine-tuning T5-Base with different methods.
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 weight decay  $= 0$ . 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 9, 10, 11, respectively.
F.1 Computational Infrastructure
The computational infrastructure information is given in Table 8.
Table 8: 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 9: Learning rate.
Model	Method	CoLA	SST-2	MRPC	QQP	MNLI	QNLI	RTE	STS-B
BERT-Base	WTA-CRS@0.3	2e-5
LoRA+WTA-CRS@0.3	2e-4	5e-4	2e-4	3e-4	3e-4	2e-4	2e-4	3e-4
T5-Base	WTA-CRS@0.3	3e-5	3e-6	3e-5	3e-5
WTA-CRS@0.1	3e-5
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	3e-5	3e-5	3e-4	3e-4
BERT-Large	WTA-CRS@0.3	2e-5
LoRA+WTA-CRS@0.3	3e-4	2e-4	2e-4	2e-4	2e-4	2e-4	3e-4	3e-4
T5-Large	WTA-CRS@0.3	3e-5	3e-6	3e-5	3e-5
WTA-CRS@0.1	3e-5	3e-6	3e-5	3e-5
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	3e-5	3e-5	\( 3e-4 \)	3e-4
T5-3B	LoRA+WTA-CRS@0.3	3e-4	3e-5	3e-4	3e-4	3e-4	3e-5	3e-4	3e-4
LoRA+WTA-CRS@0.1	3e-4	3e-5	3e-4	3e-4	3e-4	3e-5	3e-4	3e-4
Table 10: Training epoch number.
Model	Method	CoLA	SST-2	MRPC	QQP	MNLI	QNLI	RTE	STS-B
BERT-Base	WTA-CRS@0.3	20	20	10	10	10	10	20	10
LoRA+WTA-CRS@0.3	60	20	20	20	20	20	40	40
T5-Base	WTA-CRS@0.3	40	10	20	10	10	10	50	20
WTA-CRS@0.1	40	10	20	10	10	10	50	20
LoRA+WTA-CRS@0.3	40	10	20	20	20	10	50	20
LoRA+WTA-CRS@0.1	40	10	20	20	20	10	50	20
BERT-Large	WTA-CRS@0.3	60	20	20	10	10	10	40	10
LoRA+WTA-CRS@0.3	60	20	20	20	20	20	40	40
T5-Large	WTA-CRS@0.3	20	10	20	10	10	10	40	20
WTA-CRS@0.1	20	10	20	10	10	10	40	20
LoRA+WTA-CRS@0.3	40	10	40	10	10	10	60	20
LoRA+WTA-CRS@0.1	40	10	20	10	10	10	60	20
T5-3B	LoRA+WTA-CRS@0.3	40	10	20	10	10	10	60	20
LoRA+WTA-CRS@0.1	40	10	20	10	10	10	60	20
Table 11: Training mini-batch size. Here for BERT models, we utilize 128 batch size except STS-B, where a batci size 16 is used. For T5 model, we use a 100 batch size across all datasets.
Model	Method	CoLA	SST-2	MRPC	QQP	MNLI	QNLI	RTE	STS-B
BERT-Base/Large	WTA-CRS@0.3	128	16
LoRA+WTA-CRS@0.3	128	16
T5-Base/Large/3B	WTA-CRS@0.3	100
WTA-CRS@0.1	100
LoRA+WTA-CRS@0.3	100
LoRA+WTA-CRS@0.1	100