SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\nabla_ {\mathbf {Y}} ^ {\updownarrow} ^ {\top} \mathbf {X} ^ {\updownarrow} = \sum_ {i = 1} ^ {2 N} \nabla_ {\mathbf {Y}, i} ^ {\updownarrow} ^ {\top} \mathbf {X} _ {i} ^ {\updownarrow} = \sum_ {i = 1} ^ {2 N} \nabla_ {\mathbf {W} _ {i}}, \tag {7}
$$
where  $\nabla_{\mathbf{W}_i} = \nabla_{\mathbf{Y},:i}^{\ddagger}\mathbf{X}_i^{\ddagger}$ . Due to the sparsity of  $\nabla_{\mathbf{Y}}$ , the matrices  $\nabla_{\mathbf{W}_i}$  differ in magnitude and small matrices can be discarded without having a big influence on the result.
Our proposed LSS assigns each  $\nabla_{\mathbf{W}_i}$  a probability  $p_i \in [0,1], i = 1,\dots,2N$ , that satisfies  $\sum_{i=1}^{2N} p_i = N$ . We define random masks  $m_i \sim \mathrm{Bern}(p_i)$  and mask matrix  $\tilde{\mathbf{M}}$ , and approximate it as
$$
\nabla_ {\mathbf {Y}} ^ {\updownarrow} ^ {\top} \mathbf {X} ^ {\updownarrow} \approx \nabla_ {\mathbf {Y}} ^ {\updownarrow} ^ {\top} \tilde {\mathbf {M}} \mathbf {X} ^ {\updownarrow} = \sum_ {i = 1} ^ {2 N} \frac {m _ {i}}{p _ {i}} \nabla_ {\mathbf {Y};, i} ^ {\updownarrow} ^ {\top} \mathbf {X} _ {i} ^ {\updownarrow}, \text {w h e r e} \tilde {\mathbf {M}} = \operatorname {d i a g} \left(\frac {m _ {1}}{p _ {1}}, \dots , \frac {m _ {2 N}}{p _ {2 N}}\right),
$$
which is an unbiased approximation since  $\mathbb{E}\left[\nabla_{\mathbf{Y}}^{\updownarrow}{}^{\top}\tilde{\mathbf{M}}\mathbf{X}^{\updownarrow}\right] = \nabla_{\mathbf{Y}}^{\updownarrow}{}^{\top}\mathbb{E}\left[\tilde{\mathbf{M}}\right]\mathbf{X}^{\updownarrow} = \nabla_{\mathbf{Y}}^{\updownarrow}{}^{\top}\mathbf{X}^{\updownarrow}$ .
In expectation, there are only  $N$  nonzero  $m_{i}$ s. Therefore, LSS reduces the cost of MM by half. For LSS to be accurate, we minimize its variance. We have:
Proposition 4.1. (LSS variance for weight gradient)
$$
\operatorname {V a r} \left[ \sum_ {i = 1} ^ {2 N} \frac {m _ {i}}{p _ {i}} \nabla_ {\mathbf {Y} _ {i,: i}} ^ {\updownarrow} \mathbf {X} _ {i} ^ {\updownarrow} \right] = \sum_ {i = 1} ^ {2 N} \frac {1 - p _ {i}}{p _ {i}} \| \nabla_ {\mathbf {Y} _ {i,: i}} ^ {\updownarrow} \| ^ {2} \| \mathbf {X} _ {i,: i} ^ {\updownarrow} \| ^ {2}, w h e r e \operatorname {V a r} \left[ \mathbf {X} \right] := \mathbb {E} \left[ \| \mathbf {X} - \mathbb {E} \mathbf {X} \| \right] _ {F} ^ {2}.
$$
The coefficient  $c_{i} \coloneqq \| \nabla_{\mathbf{Y}_{i,:}}^{\updownarrow} \| \| \mathbf{X}_{i,:}^{\updownarrow} \|$  is called the leverage score, which can be easily computed in low time complexity. When  $p_{i} \propto c_{i}$ , the variance attends its minimum due to Cauchy inequality:
$$
\sum_ {i = 1} ^ {2 N} \frac {1}{p _ {i}} \| \nabla_ {\mathbf {Y} _ {i,:}} ^ {\ddagger} \| ^ {2} \| \mathbf {X} _ {i,:} ^ {\ddagger} \| ^ {2} = \sum_ {i = 1} ^ {2 N} \frac {c _ {i} ^ {2}}{p _ {i}} = \sum_ {i = 1} ^ {2 N} \frac {c _ {i} ^ {2}}{p _ {i}} \sum_ {i = 1} ^ {2 N} p _ {i} \geq (\sum_ {i = 1} ^ {2 N} c _ {i}) ^ {2},
$$
where the equality holds when  $p_i \propto c_i$ . Intuitively, LSS can approximate the MM Eq. (7) well with significantly lower computational cost when the leverage scores  $\{c_i\}$  are diverse, which is indeed the case as shown in Fig. 2.
Define  $\mathbf{M}^{\dagger}$  to be the top-left  $N\times N$  submatrix of  $\tilde{\mathbf{M}}$  and  $\mathbf{M}^{\downarrow}$  to be the bottom-right one, we have
$$
\nabla_ {\mathbf {Y}} ^ {\uparrow} \mathbf {\Lambda} ^ {\top} \tilde {\mathbf {M}} \mathbf {X} ^ {\uparrow} = s _ {\uparrow} \nabla_ {\mathbf {Y}} ^ {\uparrow} \mathbf {\Lambda} ^ {\top} \tilde {\mathbf {M}} ^ {\uparrow} \hat {\mathbf {X}} + s _ {\downarrow} \nabla_ {\mathbf {Y}} ^ {\downarrow} \mathbf {\Lambda} ^ {\top} \tilde {\mathbf {M}} ^ {\downarrow} \hat {\mathbf {X}},
$$
Table 1: Results on language model fine-tuning, transformer pretraining, and vision transformers fine-tuning and pretraining. Standard deviation is reported as subscript. FT refers to Fine-tuning, and PT refers to Pre-training. For WMT the result of 25.4 is result of Ultra-Low, not INT8.
DATASET	TRAIN TYPE	MODEL	METRIC NAME	BASESLINESE	4-BIT TRAINING METHODS
FP	INT8	LSQ+LUQ	HQ+LSS
GLUE-DEV	FT	BERT-BASE	Avg	82.670.24	81.450.13	75.290.52	80.810.31
BERT-LARGE	Avg	84.570.42	82.740.24	55.932.47	82.250.58
SQUAD v1	FT	BERT-BASE	F1	88.320.30	88.420.20	85.750.31	87.600.25
SQUAD v2	FT	BERT-BASE	F1	76.040.68	75.630.07	71.020.41	74.630.18
ADVERSARIAL QA	FT	BERT-BASE	F1	40.990.38	40.170.58	31.850.30	38.700.77
SWAG	FT	BERT-BASE	ACC	79.840.10	79.180.19	70.791.20	77.490.16
CONLL	FT	BERT-BASE	ACC	93.380.08	93.130.14	87.630.39	91.900.48
WMT	PT	TRANSFORMER-BASE	BLEU	27.5	25.4(ULTRA LOW)	27.17	-
SACREBLEU	26.5	-	-	25.57
CIFAR10	FT	ViT-B/32	TOP1 ACC	98.770.03	98.590.02	97.760.10	98.360.05
ViT-L/32	98.98	98.76	98.38	98.47
CIFAR100	FT	ViT-B/32	TOP1 ACC	91.940.11	90.990.07	88.630.085	89.780.06
ViT-L/32	93.07	92.2	90.97	91.13
IMAGENET1K	FT	ViT-B/32	TOP1 ACC	81.88	80.42	77.25	79.18
ViT-L/32	81.62	81.3	77.41	80.06
ViT-L/16	84.55	83.05	82.4	82.61
PT	DEiT-SMALL	TOP1 ACC	73.1	70.95	69.96	69.18
which can be implemented by two INT4 MMs with sampled rows/columns. Putting everything together, we propose the following MM procedure to compute the weight gradient:
Procedure LSS-MM
1. Quantize  $\nabla_{\mathbf{Y}}$  with BS to obtain  $\nabla_{\mathbf{Y}}^{\uparrow}$  and  $\nabla_{\mathbf{Y}}^{\downarrow}$  in INT4.
2. Compute the leverage score  $\| \nabla_{\mathbf{Y}_{i,:}}^{\ddagger}\| \| \mathbf{X}_{i,:}^{\ddagger}\|$  in FP16.
3. Sample the masks  $\{m_i\}$ .
4. Sample rows of  $\nabla_{\mathbf{Y}}$  and  $\bar{\mathbf{X}}$  given the masks  $\{m_i\}$ .
5. Compute INT4 MMs  $\nabla_{\mathbf{Y}}^{\uparrow}{}^{\top}\tilde{\mathbf{M}}^{\uparrow}\hat{\mathbf{X}}$  and  $\nabla_{\mathbf{Y}}^{\downarrow}{}^{\top}\tilde{\mathbf{M}}^{\downarrow}\hat{\mathbf{X}}$
6. Dequantize and sum up the resultant INT32 matrices to obtain the FP16 result  $\nabla_{\mathbf{Y}}^{\ddagger} \tilde{\mathbf{M}} \mathbf{X}^{\dagger}$ .
As  $\tilde{\mathbf{M}}$  only has  $N$  non-zero elements in expectation, the two matrix multiplications in Step 5 take about  $2NCD$  INT4 MACs, which aligns with the cost of the naive MM  $s_{\uparrow}\nabla_{\mathbf{Y}}^{\uparrow}\hat{\mathbf{X}}$ . The overhead of all the other steps is  $O(NC + ND)$  in total.
Activation Gradient Similar to the previous discussion, the gradient of input can be written as
$$
\nabla_ {\mathbf {Y}} \hat {\mathbf {W}} \approx \left(s _ {\uparrow} \nabla_ {\mathbf {Y}} ^ {\uparrow} + s _ {\downarrow} \nabla_ {\mathbf {Y}} ^ {\downarrow}\right) \hat {\mathbf {W}} = s _ {\uparrow} \nabla_ {\mathbf {Y}} ^ {\uparrow} \hat {\mathbf {W}} + s _ {\downarrow} \nabla_ {\mathbf {Y}} ^ {\downarrow} \hat {\mathbf {W}} = \left(\hat {\mathbf {I}} ^ {\uparrow} \nabla_ {\mathbf {Y}} ^ {\updownarrow}\right) \hat {\mathbf {W}}, \tag {8}
$$
where we define  $\nabla_{\mathbf{Y}}^{\updownarrow} = [s_{\uparrow}\nabla_{\mathbf{Y}}^{\uparrow};s_{\downarrow}\nabla_{\mathbf{Y}}^{\downarrow}]\in \mathbb{R}^{2N\times C}$  and  $\hat{\mathbf{l}}^{\ddagger} = [\mathbf{I}\quad \mathbf{I}]$  to be a  $N\times 2N$  INT4 matrix,  $\mathbf{I}$  is a  $N\times N$  identity matrix. The original product can also be implemented by two INT4 MMs  $\nabla_{\mathbf{Y}}^{\uparrow}\hat{\mathbf{W}}$  and  $\nabla_{\mathbf{Y}}^{\downarrow}\hat{\mathbf{W}}$ . But different from weight gradients, we now focus on  $\hat{\mathbf{l}}^{\updownarrow}\nabla_{\mathbf{Y}}^{\updownarrow}$  in Eq. (8) and do leverage score sampling on this MM. A detailed discussion can be found in Appendix B.2, and we only present the leverage score here. Similarly, we write the MM as the sum of  $2N$  smaller multiplications:
$$
\hat {\mathbf {I}} ^ {\updownarrow} \nabla_ {\mathbf {Y}} ^ {\updownarrow} = \sum_ {i = 1} ^ {2 N} \hat {\mathbf {I}} _ {:, i} ^ {\updownarrow} \nabla_ {\mathbf {Y} i} ^ {\updownarrow} \approx \frac {m _ {i}}{p _ {i}} \sum_ {i = 1} ^ {2 N} \nabla_ {\mathbf {Y} _ {i}},
$$
where we define  $\nabla_{\mathbf{Y}_i} = \hat{\mathbf{I}}_{;i}^{\updownarrow}\nabla_{\mathbf{Y}i}^{\updownarrow}$  and associate the probability  $p_i$  and Bernoulli mask  $m_{i}\sim \mathrm{Bern}(p_{i})$  with the  $i$  multiplication. The leverage score for activation gradient is  $c_{i}\coloneqq \| \nabla_{\mathbf{Y}i}^{\updownarrow}\|$ , and the variance attains minimum when  $p_i\propto c_i$ . More details about the algorithm can be found at Appendix. A.3 On the implementation side, once the mask  $\{m_i\}$  is known, we can decompose the MM Eq. (8) as two INT4 MMs:  $\left(\hat{\mathbf{I}}^{\updownarrow}\tilde{\mathbf{M}}\nabla_{\mathbf{Y}}^{\updownarrow}\right)\hat{\mathbf{W}} = s_{\uparrow}\tilde{\mathbf{M}}^{\uparrow}\nabla_{\mathbf{Y}}^{\uparrow}\hat{\mathbf{W}} +s_{\downarrow}\tilde{\mathbf{M}}^{\downarrow}\nabla_{\mathbf{Y}}^{\downarrow}\hat{\mathbf{W}}.$
We evaluate our INT4 training algorithm on a wide variety of tasks including language model fine-tuning, machine translation, and image classification. We implement our proposed HQ-MM and
LSS-MM algorithms with CUDA and cutlass $\dagger$ , and the implementation details can be found in Appendix A. We replace all the floating-point linear operators with our INT4 implementation except simply using LSQ for embedding layers, and leaving the last classifier layer in full precision. We adopt default architectures, optimizers, schedulers, and hyper-parameters for all the evaluated models.
5.1 Converged Model Accuracy
We compare the accuracy of the converged model on various tasks in Table 1. The compared methods include full-precision training (FP), INT8 training [3](INT8), FP4 training [47] ("Ultra-low"), 4-bit logarithm quantization [8] with LSQ for activations and weights (LSQ+LUQ), and our algorithm which utilizes HQ for forward and LSS for backpropagation (HQ+LSS). Ultra-low does not have a public implementation, so we only report its performance from its original paper on the machine translation task. Except for the large machine translation task and the task of large vision transformers, we repeat each run by three times and report the standard deviation as subscripts in tables. We do not include any kind of knowledge distillation or data augmentation.
Language model fine-tuning: We use the pretrained BERT-base-uncased and BERT-large-uncased [24] model, and evaluate the performance of our method on GLUE dev-set [53], SQUAD [41], SQUADv2 [40], Adversarial QA [4], CoNLL-2003 [42] and SWAG [61] datasets. We present the average result of the bert-base-uncased and bert-large-uncased models on the GLUE dataset. The full results are listed in Appendix C.2. Compared with LSQ+LUQ, our method achieves  $5.5\%$  improvement of accuracy on average for the bert-base model and achieves  $>25\%$  improvement of accuracy on average for the bert-large model. We further show the result on the SQUAD, SQUAD 2.0, Adversarial QA, CoNLL-2003, and SWAG datasets. On all of the tasks, compared with LSQ+LUQ, our method achieves better performance. We improve by  $1.8\%$  and  $3.6\%$  on SQUAD and SQUAD 2.0 compared to LSQ+LUQ, respectively. On the more difficult Adversarial QA, we improve by  $6.8\%$  on F1 score. On SWAG we improve by  $6.7\%$  and on CoNLL-2003 we improve by  $4.2\%$  accuracy.
Machine translation: We also apply our method for pretraining. We train a Transformer-base [52] model on WMT 14 En-De dataset [6] for machine translation. Note that we reproduce this experiment with Fairseq's recipe  $\ddagger$ , which reports the SacreBleu score (26.5 for FP) [37], while Ultra-low and LUQ report the more optimistic original BLEU score (27.5 for FP) [36]. Our HQ+LSS has about  $1.0\%$  BLEU degradation, which is smaller than  $2.1\%$  of Ultra-low and higher than  $0.3\%$  reported in the LUQ paper. Nevertheless, HQ+LSS still performs comparably with existing methods for this pretraining task, and it supports contemporary hardware.
Image Classification: We load ViT checkpoints pretrained on ImageNet21k [13], and fine-tune it on CIFAR-10, CIFAR-100 [28], and ImageNet1k. We use ViT-B/32 and ViT-L/32 for CIFAR datasets and use ViT-B/32, ViT-L/32 and ViT-L/16 for ImageNet1k. On CIFAR10 we achieve  $< 0.5\%$  accuracy degradation, while LSQ+LUQ has  $1\%$  degradation for ViT-B/32 and  $0.6\%$  degradation for ViT-L/32. On CIFAR100, INT8 already has  $\sim 1\%$  accuracy degradation, which shows its difficulty. We improve by  $1.1\%$  accuracy for ViT-B/32 and  $0.2\%$  accuracy for ViT-L/32 compared with LSQ+LUQ. On ImageNet1k, we improve by  $2\%$  accuracy for ViT-B/32,  $2.6\%$  accuracy for ViT-L/32 and  $0.2\%$  for ViT-L/32 compared with LSQ+LUQ. We further test the effectiveness of our algorithm for pretraining a DeiT-Small model [51] on ImageNet1K, where HQ+LSS can still converge to similar accuracy level compared to LSQ+LUQ, while being more hardware friendly.
Here, we conduct ablation studies to show the effectiveness of our forward and backward methods independently on the challenging CoLA dataset. To study the effectiveness of different quantizers for forward propagation, we leave backpropagation in FP16. The result is shown in Fig. 3(a). We first validate the claim in Sec. 3.2 that outliers are the main cause of accuracy degradation in quantized forward propagation. We test an "outlier" method which maintains  $1\%$  largest activation entries in FP. The "outlier" method achieves good performance, which proves that outliers are indeed the most
$\dagger$ https://github.com/NVIDIA/cutlass
$\ddagger$ https://github.com/facebookresearch/fairseq
Figure 3: CoLA performance under different methods using different bits. (a) Comparison of forward methods. (b) Comparison of backward methods.
Figure 4: Comparison of basic FP16 MM, HQ, and LSS operators.
Figure 5: SpeedUp of our INT4 training algorithm compared with FP16 PyTorch AMP on (a) Bert-Large (b) Gpt2-base.
significant challenge of the transformer's forward quantization. The hardware-unfriendly "outlier" method serves as an upper bound of methods to handle outliers. Our HQ outperforms LSQ by better handling the outliers and achieves comparable results to maintaining the outliers.
We also investigated whether more granular quantizers, such as per-token quantization or per-channel quantization could be used to quantify outliers, or whether existing methods like SmoothQuant [58] could be used for INT4 FQT. The results are listed in Appendix C.3, and we find that without HQ, none of these methods achieve good accuracy under 4-bit quantization, and the result of HQ is not strongly affected when more granular quantization methods are applied.
For backpropagation, we compare a simple minimax quantizer [3], LUQ [8] and our LSS, and leave forward propagation in FP16. The minimax quantizer divides the numerical range from the minimum to the maximum into equally large quantization bins. The result is shown in Fig. 3(b). While the bit-width is higher than 2, our LSS achieves results that are comparable and even slightly higher than LUQ. Meanwhile, LSS is more hardware friendly as it requires only INT4 arithmetic.
5.3 Computational and Memory Efficiency
Finally, we demonstrate the potential of our method to accelerate neural network training by evaluating our prototypical implementation discussed in Appendix A.6. We emphasize that our implementation is not fully optimized. For example, the backward computation requires an INT4 MM in the form of  $\mathbf{Y} = \mathbf{AB}$ , while cutlass only supports  $\mathbf{Y} = \mathbf{AB}^{\top}$ , so explicit transpose is required. We also do not fuse the linear operators with nonlinearities and normalizations. Therefore, the results cannot fully reflect the potential of INT4 training algorithms. A fully optimized implementation requires heavy engineering, which exceeds the scope of our paper.
Operator Speed: We compare the throughput of our proposed HQ-MM (HQ), LSS for computing weight gradient (LSSWeight), LSS for computing activation gradient (LSSAct), and their average throughput (INT4) with a baseline tensor-core FP16 GEMM implementation (FP16) provided by cutlass in Fig. 4 on an Nvidia RTX 3090 GPU which has a peak throughput at 142 FP16 TFLOPs and 568 INT4 TFLOPs. As the matrix size grows, the overhead of quantization diminishes and our INT4 operators can be up to 2.2 times faster compared with FP16 MM. We further analyze the quantization overhead for each operator in Appendix C.5.
Training Throughput: We compare the training throughput of the FP16 PyTorch AMP and our INT4 training algorithm for training BERT [24] and GPT [38]-style language models on a system of 8 Nvidia A100 GPUs. We vary the hidden layer size, intermediate fully-connected layer size, and batch size, and plot the speedup of INT4 training in Fig. 5. Our INT4 training algorithm can achieve up to  $35.1\%$  speedup and an average of  $15.8\%$  for BERT-style models and up to  $26.5\%$  speedup and an average of  $19.2\%$  for GPT-style models. The training time can be found in Appendix C.4.
Inference Speed: We compare the inference speed of our algorithm with I-BERT [25] by comparing the speedup numbers reported in its original paper. Following the I-BERT paper, we compare the speedup of integer-based inference algorithms relative to a FP32 baseline on an Nvidia T4 GPU on the BERT-base and BERT-large models, and test with sequence lengths of 128 and 256. While I-BERT only reported speedup numbers for batch sizes 1, 2, 4, 8, we test the speedup for batch sizes
Table 2: The inference speed up compared with I-BERT. Since the open-sourced versions of I-BERT are based on fake quantization, we compare the inference speed of our algorithm with I-BERT by comparing the speedup numbers reported in its original paper with respect to FP32.
SL	MODEL	METHOD	BATCH SIZE
1	2	4	8	16	32	64	128	256	512	1024
128	BERT-BASE	I-BERT	2.42	3.36	3.39	3.31	-	-	-	-	-	-	-
Ours	0.17	0.25	0.46	0.87	1.22	2.07	3.61	3.57	4.5	4.54	4.92
BERT-LARGE	I-BERT	3.2	4	3.98	3.81	-	-	-	-	-	-	-
Ours	0.22	0.43	0.81	1.46	2.13	3.34	4.28	4.81	5.4	6.08	6.48
256	BERT-BASE	I-BERT	3.11	2.96	2.94	3.15	-	-	-	-	-	-	-
Ours	0.25	0.54	0.95	1.53	1.94	3.2	3.76	4	4.15	4.15	4.14
BERT-LARGE	I-BERT	3.19	3.51	3.37	3.4	-	-	-	-	-	-	-
Ours	0.45	0.86	1.47	2.44	2.67	3.99	4.87	5.24	5.11	5.41	OOM
ranging from 1 to 1024 to better reflect the performance of throughput-oriented scenarios (such as a cloud language model service provider).
In Table 2 we report the speed up result. While I-BERT's speedup numbers seems to be insensitive to the batch size and sequence length, our speedup increases with the batch size. I-BERT shows up to  $3.98\mathrm{x}$  speedup for smaller batch size, while our algorithm can achieve higher speedup for batch sizes higher than 64, and eventually gives a speedup of  $6.48\mathrm{x}$  for a sequence length of 128 and a batch size of 1024.
For the BERT-base model, our method shows an inference speed improvement of 3.57-4.92 times compared to FP32 for batch sizes larger than 64. In comparison, I-BERT achieved an inference speed improvement of 2.42-3.39 times compared to FP32 when the batch size was small. For BERT-large, our method shows an inference speed improvement of 4.81-6.48 times compared to FP32 for batch sizes larger than 64. In comparison, I-BERT achieved an inference speed improvement of 3.20-4.00 times compared to FP32 when the batch size is small. Therefore, our algorithm can potentially achieve higher throughput than I-BERT.
We propose a hardware-friendly INT4 training method for transformers. By analyzing the properties of MMs in transformers, we propose HQ and LSS methods to quantize activations and gradients while preserving accuracy. On several important tasks, our method performs comparably or better than existing INT4 methods. Our work can be potentially extended beyond transformers to other MM-only architectures, such as MLP-Mixer [50], graph neural networks [26], and recurrent neural networks [20]. We leave it as a future direction.
Broader Impacts: Our algorithm can improve efficiency and reduce the energy consumption of training neural networks, which helps reduce the carbon footprint caused by deep learning. However, our efficient training algorithm might also facilitate the development of large language models with safety concerns for human beings; and malicious AI applications such as fake content generation.
**Limitations:** Our INT4 training algorithm does not support convolutional layers. Moreover, our algorithm cannot yet work well for those extremely large models such as OPT-175B. LSS utilizes gradient sparsity, and may not work well for certainty pretraining tasks. Finally, our algorithm may incur more memory accesses than FP16 training algorithms, which might affect the speedup.
The authors would like to thank Weilin Zhao, Sixuan Ma, Xiaoxuan Liu, Ziteng Wang, Bingrui Li, Cheng Lu, and Zhiyuan Liu for valuable discussions and help on the training of large-scale language models. This work was supported by the National Key Research and Development Program of China (No. 2021ZD0110502), NSFC Projects (Nos. 62061136001, 62106123, 62076147, U19A2081, 61972224, 62106120), Tsinghua Institute for Guo Qiang, and the High Performance Computing Center, Tsinghua University. J.Z is also supported by the XPlorer Prize.
In this section, we present some works that need to be done to actually accelerate the training process on hardware.
A.1 BMM in Attention
In attention, there are batch matrix multiplications (BMMs) that need to be dealt with. We now show that our method for MMs can be extended to BMMs.
Consider the following BMM product:
$$
\mathbf {T} = \operatorname {B M M} (\mathbf {Q}, \mathbf {K} ^ {\top}),
$$
where we define  $\mathbf{T} \in \mathbb{R}^{B \times N \times P}$ ,  $\mathbf{Q} \in \mathbb{R}^{B \times N \times M}$ ,  $\mathbf{K} \in \mathbb{R}^{B \times P \times M}$ . The Hadamard matrix is defined as:
$$
\hat {\mathbf {H}} = \operatorname {R e p a t} _ {B} (\mathbf {H}) = \operatorname {R e p a t} _ {B} (\operatorname {B l o c k D i a g} (\mathbf {H} _ {k}, \dots , \mathbf {H} _ {k})),
$$
where  $\hat{\mathbf{H}}\in \mathbb{R}^{B\times M\times M},\mathbf{H}\in \mathbb{R}^{M\times M},\mathbf{H}_k\in \mathbb{R}^{2^k\times 2^k}$ . In this case,
$$
\mathbf {T} \approx \operatorname {B M M} \left(\operatorname {B M M} (\mathbf {Q}, \hat {\mathbf {H}}), \operatorname {B M M} (\mathbf {K}, \hat {\mathbf {H}}) ^ {\top}\right),
$$
which verifies that our HQ can be applied to BMMs.
For backward, the gradient of weight and activation can be calculated by the straight-through estimator  $\lfloor x\rfloor^{\prime} = 1$  and the chain rule:
$$
\begin{array}{l} \nabla_ {\mathbf {Q}} = s _ {Q} \left(\mathrm {B M M} (\nabla_ {\mathbf {T}} ^ {\top}, \hat {\mathbf {K}}) \circ \mathbb {I} _ {Q}\right) \mathbf {H} ^ {\top}, \\ \nabla_ {\mathbf {K}} = s _ {K} \mathbb {I} _ {K} \circ \operatorname {B M M} (\nabla_ {\mathbf {T}}, \hat {\mathbf {Q}}) \mathbf {H} ^ {\top} = s _ {K} \operatorname {B M M} (\mathbb {I} _ {K} \circ \nabla_ {\mathbf {T}}, \hat {\mathbf {Q}}) \mathbf {H} ^ {\top}, \\ \end{array}
$$
where we define  $s_Q \in \mathbb{R}^B, s_k \in \mathbb{R}^B$  being the batch step size,  $\hat{\mathbf{K}} = \mathrm{int}_{s_K}\left(\mathrm{BMM}(\mathbf{K},\hat{\mathbf{H}})\right)$ ,  $\hat{\mathbf{Q}} = \mathrm{int}_{s_Q}\left(\mathrm{BMM}(\mathbf{Q},\hat{\mathbf{H}})\right)$ ,  $\mathbb{I}_Q = \mathbb{I}(-Q_N \leq \mathbf{Q} / s_Q \leq Q_P)$ , and  $\mathbb{I}_K = \mathbb{I}(-Q_N \leq \mathbf{K} / s_K \leq Q_P)$ .
Similar to Sec. 4.2, we only focus on  $\mathrm{BMM}(\nabla_{\mathbf{T}}^{\top},\hat{\mathbf{K}})$  and  $\nabla_{\mathbf{T}}$  , since we do leverage sampling on them.
For  $\mathrm{BMM}(\nabla_{\mathbf{T}}^{\top},\hat{\mathbf{K}})$ , we define the sample probability  $p_i$  and sample the  $\tilde{\mathbf{M}}$  in the same way as MMs. The matrix can be computed as  $\mathrm{BMM}(\mathrm{BMM}(\nabla_{\mathbf{T}}^{\top},\hat{\tilde{\mathbf{H}}}),\hat{\mathbf{K}}^{\top})$ , where  $\hat{\tilde{\mathbf{H}}}$  is defined as  $\mathrm{CONCAT}(\tilde{\mathbf{H}}_1,\dots ,\tilde{\mathbf{H}}_B)$ ,  $\nabla_{\mathbf{T}}^{\uparrow \top}$  and  $\hat{\mathbf{K}}^{\uparrow}$  follows the same definition of Eq. 6 and the leverage score is  $c_{b,i} := \| \nabla_{\mathbf{T}b,i,:}^{\uparrow}\| \| \mathbf{K}_{b,i,:}^{\uparrow}\|$  for  $0\leq b\leq B,0\leq i\leq 2M$ .