SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
C E _ {n} ^ {T} = - \sum_ {i = 1} ^ {V} y _ {n, i} \log \left(\mathbf {P} _ {n, i} ^ {T}\right), \quad \operatorname {s c a l e} _ {n} = \frac {e ^ {\mathrm {C E} _ {n} ^ {T} / \tau}}{\sum_ {k = 1} ^ {N} e ^ {\mathrm {C E} _ {k} ^ {T} / \tau}} \tag {6}
$$
$$
L _ {T S L D} = \sum_ {n = 1} ^ {N} \left(\operatorname {s c a l e} _ {n} \times - \sum_ {i = 1} ^ {V} \mathbf {P} _ {n, i} ^ {T} \log \left(\mathbf {P} _ {n, i} ^ {S}\right)\right) \tag {7}
$$
The implementation of TSLD is straightforward. By considering the relationship between token confidence and token prediction cross entropy loss in Fig 3(b), we can determine the  $n_{th}$  token scale values  $(\mathrm{scale}_n)$  based on the cross entropy loss of the teacher model  $(CE_n^T)$  using softmax function as shown in Eq. 6. Note that  $y_{n,i}$  is a ground truth label where  $y_{n,i} = 1$  if token  $i$  is the true next token at position  $n$  and  $\tau$  is the temperature parameter for softmax function. The scale of the each token  $(\mathrm{scale}_n)$  is then applied in the logit distillation by multiplying it with the cross entropy loss between the student and teacher models as shown in Eq. 7. As depicted in Fig. 3(c) left, the scale values for token-specific Logit KD are determined adaptively based on the per-token cross-entropy loss of the teacher model. In the right graph of Fig. 3(c), we can observe that tokens with higher cross entropy loss in the teacher model correspond to higher scale values in the per token scale graph, compared to Logit KD method applying the same scale value  $(1 / N)$  across all tokens as shown in the left graph of Fig. 3(c).
TSLD brings about two significant effects by applying different scales based on confidence disparity, with negligible computational overhead. As shown in Fig. 3(a), TSLD de-emphasizes Logit KD for high-confidence tokens, thereby preventing overfitting. Conversely, for low-confidence tokens possessing a high entropy probability distribution, TSLD emphasizes the Logit KD. This action allows the student model to more closely mimic the teacher model's cross-entropy loss as seen in the right graph of Fig. 3(c). A detailed analysis of the TSLD method's computational cost can be found in Appendix A.1.
In this section, we evaluate the effectiveness of TSLD in the QAT of various sizes of decoder models with sub-4bit quantization. We've set up comparative experiments to demonstrate the proficiency of our TSLD method against existing PTQ and other QAT KD methods. Our findings illustrate that TSLD substantially enhances both the language modeling performance (measured by Perplexity or PPL) and the accuracy in tasks related to reasoning (common-sense QA and arithmetic) and natural language understanding.
Task and Models. We evaluate our proposed method for language modeling (PTB [34]), commonsense QA tasks (PIQA [35], OpenbookQA [36], ARC_easy [37], ARC_challenge [37] and arithmetic
Precision	Quantization Method	Optimization Method	GPT-2	OPT
0.1B	0.3B	0.8B	1.5B	0.1B	1.3B	2.7B	6.7B
FP16 baseline	20.91	18.21	15.20	14.26	18.17	13.75	11.43	10.21
W4A16	PTQ	OPTQ [14]	22.41	19.35	17.26	15.86	19.75	14.30	11.82	11.73
QAT	Logit [24]	20.98	18.54	16.79	15.42	17.60	13.73	11.82	11.20
Logit+GT	21.51	18.58	15.49	14.89	19.63	15.03	12.58	11.78
TSLD	19.95	17.53	15.32	14.50	17.45	13.90	11.59	11.00
W2A16	QAT	L2L+Logit [25]	23.79	21.21	17.80	15.82	20.47	17.62	14.67	11.75
Logit [24]	22.84	19.87	16.46	15.27	18.86	14.80	12.26	11.33
Logit+GT	23.80	20.20	17.77	16.52	21.62	16.41	13.20	12.41
TSLD	21.74	18.57	16.14	15.02	18.58	14.60	11.97	11.17
Table 1: Perplexity comparison in GPT-2 and OPT series across various model sizes (0.1B to 6.7B) on the PTB dataset with QAT-KD (tensor-wise) and PTQ (channel-wise) quantization methods
QAT KD
Method	PIQA	OpenbookQA	ARCeasy	ARC_challenge	GSM8K
ACC (↑)	PPL (↓)	ACC (↑)	PPL (↓)	ACC (↑)	PPL (↓)	ACC (↑)	PPL (↓)	ACC (↑)	PPL (↓)
OPT-2.7B FP16	76.71	10.91	49.60	26.16	66.12	7.41	37.20	8.96	20.39	2.07
Logit [24]	74.32	11.69	45.40	29.41	58.92	9.05	31.91	12.38	20.02	2.03
GT+Logit	74.97	12.10	46.20	31.08	58.84	8.66	32.16	12.04	19.56	2.12
TSLD	75.62	11.35	46.81	28.93	59.39	8.12	33.45	11.05	20.24	2.03
Model	GPT-Neo-1.3B	OPT-6.7B	LLaMA-7B
QAT KD	PTB (PPL)	GSM8K (ACC/PPL)	PTB (PPL)	GSM8K (ACC/PPL)
FP16	17.62 (↓)	22.52 (↑)	1.89 (↓)	8.76 (↓)	30.25 (↑)	1.47 (↓)
Logit [24]	21.01	21.08	1.93	12.22	25.47	1.52
TSLD	19.27	24.49	2.14	11.60	26.23	1.52
Table 2: Top: Results for the OPT-2.7B model fine-tuned on common-sense QA and arithmetic reasoning task using various QAT-KD (tensor-wise) methods. Bottom: QAT-KD (channel-wise) results on language modeling task and arithmetic reasoning task across various GLM models.
reasoning based text-generation task (GSM8K [38]). Additionally, our assessment extends to Natural Language Understanding (NLU) task (GLUE [39]), ensuring a comprehensive analysis. Our benchmark models encompass widely used GLMs, such as GPT-2 [2], OPT [4], GPT-Neo [40] and LLaMA [5] [41] with various sizes ranging from 0.1B to 7B parameters.
Fine-Tuning Settings. In fine-tuning the language modeling task, we employ a chunk-based preprocessing method: all training datasets are concatenated, then split into shorter chunks defined by input sequence length. For reasoning task fine-tuning, we utilize a sentence-based approach, concatenating each dataset's question and answer parts to form new sentences, individually used as the fine-tuning dataset. Detailed hyper-parameter settings and other specifics are in Appendix C.2. Experiments are conducted on an A100-40GB GPU. Our QAT experiments start with models that have undergone task-specific fine-tuning. During quantization, the KD process employs the FP fine-tuned model as the teacher model, while the quantized model acts as the student.
Implementation Settings. We devise a QAT-KD framework that leverages pipeline parallelism with PyTorch Pipe API enabling the training of models with capacities exceeding 1.3 billion. We apply weight quantization to the matrix multiplication layers in each decoder layer of the GLM. We conduct experiments on L2L KD [25] but encounter out-of-memory problems for models with more than 1.3B parameters in A100-40GB GPU. This issue is believed to arise due to the requirement for both the teacher and student models to store the outputs generated by all of their respective intermediate layers during the knowledge distillation process. GPU memory consumption comparisons for each QAT-KD method can be found in the Appendix A1.
5.2 Evaluation on Language Modeling Task
Table 1 outlines the performance comparison of TSLD with leading PTQ and QAT methods [14, 24] for language modeling of PTB dataset. For 4-bit weight quantization, OPTQ sees a notable performance drop in GPT-2 and OPT models up to 6.7 billion parameters, aligning with the original paper's observations [14]. However, QAT methods show lower perplexity due to weight parameters fine-tuned for robust reduced-precision inference. QuantGPT [24], which exclusively uses Logit KD achieves impressive perplexity, whereas Logit+GT KD sees degradation. Conversely, TSLD offers
Precision	QAT KD Method	CoLA	MRPC	SST-2	RTE
ACC (↑)	PPL (↓)	ACC (↑)	PPL (↓)	ACC (↑)	PPL (↓)	ACC (↑)	PPL (↓)
OPT-1.3B FP16	61.03	1.34	81.92	2.58	94.26	2.00	76.53	3.94
W4A16	OPTQ [14]	54.61	1.36	80.14	2.43	95.07	2.02	56.32	3.96
AWQ [13]	13.63	1.45	66.42	3.49	94.26	2.02	54.51	4.72
Logit [24]	50.76 ±2.35	1.36	81.94 ±1.48	2.62	93.57 ±0.23	2.09	75.23 ±0.83	4.34
GT+Logit	54.07 ±0.34	1.34	83.17 ±0.51	2.60	93.34 ±0.22	2.11	75.31 ±1.07	4.09
TSLD	56.33 ±0.98	1.34	83.33 ±1.22	2.52	94.05 ±0.19	2.04	75.97 ±0.31	4.05
W2A16	Logit [24]	48.72 ±2.68	1.37	81.62 ±0.62	2.79	93.08 ±0.35	2.11	74.15 ±1.36	4.72
GT+Logit	50.10 ±1.38	1.34	82.10 ±0.99	2.65	92.77 ±0.28	2.14	73.79 ±1.16	4.44
TSLD	54.47 ±1.47	1.34	82.20 ±0.94	2.63	93.92 ±0.29	2.06	75.31 ±0.54	4.36
Table 3: Results for the OPT-1.3B model fine-tuned on GLUE [39] using different QAT-KD methods with five random seed tests. Channel-wise quantization is applied in both PTQ and QAT-KD.
the lowest perplexity, underlining token-wise scaling method's effectiveness in incorporating GT knowledge. Remarkably, TSLD's performance boost allows QAT models to match full-precision performance across various capacity ranges in all decoder models.
For 2-bit weight quantization, L2L KD sees significant accuracy degradation, and Logit+GT KD suffers from overfitting. TSLD outperforms Logit KD [24] across all model sizes, maintaining PPL degradation of no more than 1.0 from the baseline. We also tested the general applicability of TLSD on popular open-sourced GLM models (GPT-Neo-1.3B [40], LLaMA-7B [5]) in language modeling PTB task. Table 2-below indicates that TSLD consistently surpassed the competitor, Logit KD [24]. Notably, 2-bit TSLD utilizes simple ternary weight quantization, which is hardware-friendly.
5.3 Evaluation on Reasoning Task
We assess the effectiveness of our proposed method in commonsense QA (PIQA, OpenbookQA, ARCeasy, ARC_challenge) and reasoning-based text-generation task (GSM8K) employing the LM Evaluation Harness framework from EleutherAI [42]. Given the capacity requirements for reasoning tasks, we use OPT-2.7B/6.7B and LLaMA-7B models as baselines, rather than smaller models. Table 2 presents a performance comparison of 2-bit weight quantization with different KD methods. In commonsense QA tasks, TSLD consistently showcased the lowest perplexity and, consequently, the highest accuracy, drawing the 2-bit quantization results even closer to FP performance as shown in Table 2-top.
Considering the GSM8K task, Table 2 reveals that TSLD outperformed Logit KD in terms of perplexity and accuracy with OPT-2.7B/6.7B and LLaMA models. Notably, while QuantGPT (Logit KD) achieves comparable or better perplexity, its reasoning task accuracy is lower, potentially due to insufficient GT information. Conversely, TSLD achieves excellent reasoning accuracy while maintaining competitive perplexity, underscoring TSLD's ability to balance language modeling and reasoning performance through its token-wise scaling to avoid overfitting. The generated text sample results of the GSM8K task are provided in the Appendix D.
5.4 Evaluation on Language Understanding Task
We fine-tune the decoder model for Natural Language Understanding (NLU) tasks using a language modeling approach as illustrated in Fig. 1(b). In our experiments outlined in Table 3, we compare the performance of the latest PTQ methods (AWQ [13], OPTQ [14]) and QAT-KD methods with OPT-1.3B model. For 4-bit quantization, the PTQ technique shows a noticeable degradation in performance compared to QAT results, excluding SST-2 task. TSLD achieves the lowest perplexity and the highest accuracy across all the experiments except SST-2, where its accuracy is in-par with the full-precision case. These findings demonstrate that TSLD can robustify the performance of 4-bit quantized GLMs for various NLU tasks, while 4-bit PTQ may suffer from performance degradation.
In ternary quantization, TSLD consistently outperforms the alternative QAT-KD methods for all the cases, demonstrating its superior performance in bridging the accuracy gap with the full-precision cases. Interestingly, ternary TSLD even achieved similar or superior accuracy compared to 4-bit PTQ in many tasks (e.g., CoLA, MRPC, SST-2), highlighting its benefits on both accuracy and memory savings.
Model config.	6.7B	13B	175B
Input channel	4096	4096	16384	5120	5120	20480	12288	12288	49152
Output channel	4096	16384	4096	5120	20480	5120	12288	49152	12288
FP32 baseline	0.067	0.201	0.194	0.100	0.285	0.287	0.391	1.472	1.522
8-bit (× speedup)	0.039	0.068	0.068	0.043	0.095	0.092	0.121	0.407	0.395
×1.71	×2.93	×2.83	×2.32	×3.01	×3.13	×3.23	×3.61	×3.85
4-bit (× speedup)	0.030	0.057	0.057	0.041	0.076	0.075	0.096	0.326	0.318
×2.23	×3.52	×3.40	×2.45	×3.75	×3.82	×4.07	×4.51	×4.78
2-bit (× speedup)	0.025	0.053	0.053	0.039	0.072	0.064	0.077	0.221	0.232
×2.68	×3.77	×3.65	×2.57	×3.95	×4.48	×5.07	×6.66	×6.56
Reduced-Precision Kernels Execution Time. We developed custom CUDA kernels to enhance inference speed with applied (2-,4-,8-bit) quantization. Like OPTQ, we packed the weights to minimize the model size and load overhead. Our kernel eliminates the need for weight unpacking during the model forward pass, resulting in a speedup shown in Table 4. We tested our kernel mainly on models larger than 6.7B, where weight load overhead is notably high. The reported times are the average execution time for 10,000 kernel runs on a single A100-80GB GPU. For the FP32 baseline, we used PyTorch's nn.Linear. As shown in Table 4, our 2-bit kernel for the 175B model can potentially accelerate a single matrix multiplication operation by an average of approximately 6.1 times compared to FP32.
Table 4: Kernel execution time (msec)
Precision	Granularity	GPT-2	OPT	GPT-Neo 1.3B	LLaMA 7B
0.1B	0.3B	0.8B	0.1B	1.3B
FP16 baseline	20.91	18.21	15.20	18.17	13.75	17.62	8.76
W2A16	Tensor-wise	21.74	18.57	16.14	18.58	14.60	30.60	12.31
Channel-wise	21.30	18.48	15.97	18.42	14.46	19.27	11.60
Table 5: Comparison of tensor-wise and channel-wise quantization across various GLMs (GPT-2, OPT, GPT-Neo and LLaMA). The TSLD KD method is employed in this experiments
Quantization Granularity Impact. To account for output channel weight variations, channel-wise quantization is used [15, 14]. By integrating our QAT-KD approach with channel-wise quantization, we can achieve further performance enhancement. An interesting observation emerges: the gains from channel-wise quantization vary by the type of GLM. As illustrated in Table 5, for the GPT-2 and OPT series, the PPL performance increase due to channel-wise quantization is less than 1. However, for GPT-Neo and LLaMA, the performance enhancement effect resulting from channel-wise quantization is significantly pronounced. This variation in performance gains suggests distinct channel-wise weight distributions across different GLM models. A detailed analysis of the weight distribution for each GLM is addressed in the Appendix A.4.
We introduce token-scaled logit distillation, a new approach for Quantization-Aware Training of Generative Language Models. This method effectively reduces overfitting and enhances learning from the teacher model and ground truth. Importantly, this research is the first to evaluate ternary quantization-aware training on large-scale GLMs, achieving less than 1.0 perplexity degradation and preserving commonsense QA and arithmetic reasoning task accuracy.
This work was supported by Institute of Information & Communications Technology Planning & Evaluation (IITP) grants funded by the Korea government (MSIT) (2020-0-01373, Artificial Intelligence Graduate School Program Hanyang University, 2023-RS-2023-00253914, artificial intelligence semiconductor support program to nurture the best talents), the National Research Foundation of Korea (NRF) grant, funded by the Korea government (MSIT) (No. 2021R1A2C1013513, No. RS-2023-00260527), and the Artificial Intelligence Industrial Convergence Cluster Development Project, funded by the Ministry of Science and ICT (MSIT, Korea) and Gwangju Metropolitan City.
[1] Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language understanding by generative pre-training, 2018.
[2] Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners, 2019.
[3] 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.
[4] Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. Opt: Open pre-trained transformer language models. arXiv preprint arXiv:2205.01068, 2022.
[5] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971, 2023.
[6] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.02311, 2022.
[7] OpenAI. Gpt-4 technical report. arXiv preprint arXiv:2303.08774, 2023.
[8] Jason Wei, Maarten Bosma, Vincent Zhao, Kelvin Guu, Adams Wei Yu, Brian Lester, Nan Du, Andrew M. Dai, and Quoc V Le. Finetuned language models are zero-shot learners. In International Conference on Learning Representations, 2022. URL https://openreview.net/forum?id=gEzrGCozdqR.
[9] Hyung Won Chung, Le Hou, Shayne Longpre, Barret Zoph, Yi Tay, William Fedus, Yunxuan Li, Xuezhi Wang, Mostafa Dehghani, Siddhartha Brahma, Albert Webson, Shixiang Shane Gu, Zhuyun Dai, Mirac Suzgun, Xinyun Chen, Aakanksha Chowdhery, Alex Castro-Ros, Marie Pellat, Kevin Robinson, Dasha Valter, Sharan Narang, Gaurav Mishra, Adams Yu, Vincent Zhao, Yanping Huang, Andrew Dai, Hongkun Yu, Slav Petrov, Ed H. Chi, Jeff Dean, Jacob Devlin, Adam Roberts, Denny Zhou, Quoc V. Le, and Jason Wei. Scaling instruction-finetuned language models, 2022.
[10] Yizhong Wang, Yeganeh Kordi, Swaroop Mishra, Alisa Liu, Noah A Smith, Daniel Khashabi, and Hannaneh Hajishirzi. Self-instruct: Aligning language model with self generated instructions. arXiv preprint arXiv:2212.10560, 2022.
[11] Xiang Lisa Li and Percy Liang. Prefix-tuning: Optimizing continuous prompts for generation. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers), pages 4582-4597, Online, August 2021. Association for Computational Linguistics. doi: 10.18653/v1/2021.acl-long.353. URL https://aclanthology.org/2021.acl-long.353.
[12] Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto. Stanford alpaca: An instruction-following llama model. https://github.com/tatsu-lab/stanford_alpaca, 2023.
[41] Hugo Touvron, Louis Martin, Kevin R. Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajwal Bhargava, Shruti Bhosale, Daniel M. Bikel, Lukas Blecher, Cristian Cantón Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony S. Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel M. Kloumann, A. V. Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, R. Subramanian, Xia Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zhengxu Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models. ArXiv, abs/2307.09288, 2023. URL https://arxiv.org/abs/2307.09288.
[42] Leo Gao, Jonathan Tow, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Kyle McDonell, Niklas Muennighoff, Jason Phang, Laria Reynolds, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. A framework for few-shot language model evaluation, September 2021. URL https://doi.org/10.5281/zenodo.5371628.
[43] Aohan Zeng, Xiao Liu, Zhengxiao Du, Zihan Wang, Hanyu Lai, Ming Ding, Zhuoyi Yang, Yifan Xu, Wendi Zheng, Xiao Xia, Weng Lam Tam, Zixuan Ma, Yufei Xue, Jidong Zhai, Wenguang Chen, Zhiyuan Liu, Peng Zhang, Yuxiao Dong, and Jie Tang. GLM-130b: An open bilingual pre-trained model. In The Eleventh International Conference on Learning Representations, 2023. URL https://openreview.net/forum?id=-AwOrrrPUF.
[44] Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus), 2020.
[45] Abien Fred Agarap. Deep learning using rectified linear units (relu). arXiv preprint arXiv:1803.08375, 2018.
[46] Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus), 2023.
[47] Jianlin Su, Yu Lu, Shengfeng Pan, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. arXiv preprint arXiv:2104.09864, 2021.
[48] Noam Shazeer. Glu variants improve transformer, 2020.
[49] Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The Pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00027, 2020.
[50] Chaofan Tao, Lu Hou, Wei Zhang, Lifeng Shang, Xin Jiang, Qun Liu, Ping Luo, and Ngai Wong. Compression of generative pre-trained language models via quantization, 2022.
A Supplementary Analysis
A.1 Computation Requirements of TSLD
Training Method	QAT-KD Method	GPT2-0.3B 512	GPT2-0.8B 512	OPT-1.3B 256	OPT-1.3B 1024
Speed (iter/sec)	Memory (MiB)	Speed (iter/sec)	Memory (MiB)	Speed (iter/sec)	Memory (MiB)	Speed (iter/sec)	Memory (MiB)
QAT	GT	2.03	19663	1.03	36317	5.15	10051	2.83	15167
QAT-KD	Logit	1.57	22622	0.81	40989	4.44	11589	2.27	17529
GT+Logit	1.56	22622	0.81	40989	4.44	11589	2.27	17529
L2L+Logit	1.51	31462	OOM	OOM	4.28	12143	2.12	25315
TSLD	1.57	22622	0.81	40989	4.43	11589	2.26	17529
Table A1: QAT memory consumption and training speed study for KD method. The results are reported on the PTB dataset on the Ternary QAT-KD of GPT-2 and OPT series models with input sequence lengths ranging from 256 to 1024
The TSLD method integrates token-wise cross-entropy loss with Logit KD, involving two operations as detailed in Eq.7. Specifically, the term  $\sum_{i=1}^{V} y_{n,i} \log(P_{n,i}^T)$  computes the cross-entropy loss from teacher logits  $(Z_n^T)$ . This result, processed through a softmax function, derives the scaling value for each token. Multiplied element-wise with Logit KD term,  $\sum_{i=1}^{V} P_{n,i}^T \log(P_{n,i}^S)$ , it yields a token-wise scaled Logit KD. In fact, TSLD leverages the teacher logits that are pre-computed in Logit KD, circumventing extra memory usage. Furthermore, the associated computations have a complexity of  $O(N)$ , making TSLD's overhead negligible for training.
To evaluate TSLD's efficiency, we detail training speeds and GPU memory consumption for various QAT-KD methods using GPT-2 models in Table A1 left. Compared to Logit-based methods, TSLD maintains speed without extra memory consumption. In contrast, L2L KD stores intermediate activations from both the teacher and student models for KD, resulting in significantly increased memory requirements, evident from Table A1 left. As model size grows, as evidenced in scenarios utilizing GPT2-Large, memory requirements rise, leading to "Out of Memory" errors on an A100-40GB GPU. These findings highlight efficacy of TSLD, enhancing QAT-KD performance with memory comparable to Logit KD, while L2L KD demands significantly more. Even when the sequence length is extended from 256 to 1024, as Table A1 right shows, TSLD maintains the same GPU memory consumption and training speed as the Logit KD method.
A.2 Token Confidence Disparity Analysis
Our analysis of confidence disparity in token predictions, detailed in Section 4.2, extends beyond a specific GLM model. In fact, this observed trend is consistently present across various GLM models. As shown in Fig A1, we can distinctly observe the emergence of the Low Confidence Region (blue box) and the High Confidence Region (yellow box) consistently across models: OPT-6.7B(left), LLaMA-7B(middle), and LLaMA-2-7B(right). Additionally, as shown in Fig A1 right, we plot the token prediction's statistics with varying input sequence lengths of 128 and 512. Regardless of the sequence length, the demarcation of confidence disparity remains consistent. This observation demonstrates that the TSLD methodology, grounded in the the probabilistic dynamics of token prediction, can be universally applied across various GLMs.