SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Minsoo Kim $^{1}$  Sihwa Lee $^{1}$  Janghwan Lee $^{1}$  Sukjin Hong $^{1,2}$  Du-Seong Chang $^{2}$  Wonyong Sung $^{3}$  Jungwook Choi $^{1*}$
<sup>1</sup>Hanyang University, Seoul, Republic of Korea
$^{2}$ KT, Seoul, Republic of Korea
$^{3}$ Seoul National University, Seoul, Republic of Korea
{minsoo2333, macto94, hwanii0288, choij}@hanyang.ac.kr
{sukjin.hong, dschang}@kt.com, wysung@snu.ac.kr
Generative Language Models (GLMs) have shown impressive performance in tasks such as text generation, understanding, and reasoning. However, the large model size poses challenges for practical deployment. To solve this problem, Quantization-Aware Training (QAT) has become increasingly popular. However, current QAT methods for generative models have resulted in a noticeable loss of accuracy. To counteract this issue, we propose a novel knowledge distillation method specifically designed for GLMs. Our method, called token-scaled logit distillation, prevents overfitting and provides superior learning from the teacher model and ground truth. This research marks the first evaluation of ternary weight quantization-aware training of large-scale GLMs with less than 1.0 degradation in perplexity and achieves enhanced accuracy in tasks like common-sense QA and arithmetic reasoning as well as natural language understanding.[2]
Generative language models (GLMs) have made impressive strides in text generation, understanding, and reasoning, attracting significant attention in the field [1-7]. However, deploying GLMs remains a challenge due to their enormous model sizes. There is rising interest in practical GLMs with less than 10 billion parameters. Their capability can be improved through instruction fine-tuning [8-11]. For instance, Alpaca [12] showed that a fine-tuned 7 billion-parameter model can match the text generation performance of a 175 billion parameter GLM, highlighting the potential of smaller, more manageable models.
Since practical GLMs still contain billions of parameters, there is extensive research into model compression techniques for their efficient deployment. One such method is post-training quantization (PTQ), which simplifies the process by reducing bit-precision to 8 or 4 bits without the need for fine-tuning pre-trained GLMs [13-17]. This approach has gained traction due to its straightforward and fast processing time. However, it's been observed that these techniques cause a significant decrease in accuracy when the parameter count drops below 10 billion or when the bit-precision falls under 4 bits. As a result, there's a clear need for a more reliable quantization approach for GLMs with sub-4bit precision.
In response, we propose an alternative method, quantization-aware training (QAT), to address the issues PTQ poses for fine-tuned GLMs. QAT is a prevalent quantization technique that counteracts
2Our code is available at https://github.com/aiha-lab/TSLD
accuracy loss and attains a high compression rate for efficient deployment [18]. Notably, successful fine-tuning of sub-4bit natural language understanding models has been achieved through layer-to-layer (L2L) knowledge distillation (KD), a method used to offset errors resulting from aggressive quantization, such as binary or ternary weights [19-23]. However, applying QAT to GLM has limited success. While [24] introduced a token-level contrastive loss and [25] offered initial insights into the challenges of quantizing GLMs, both studies encountered substantial increases in perplexity in language modeling. Furthermore, no existing studies apply QAT to GLMs with billions of parameters, primarily due to the expensive nature of training with KD.
This study delves into the fundamental challenges of applying QAT to fine-tuned GLMs. We identify two main issues. First, the structure of a self-attention map in masked self-attention causes cumulative quantization errors across tokens, which conventional L2L KD struggles to compensate for. Second, the teacher-forcing mechanism [26] used in fine-tuning Transformer decoder necessitates ground-truth loss (GT Loss) – a factor largely overlooked in previous QAT methods – but including GT Loss risks overfitting. Our investigation reveals that logit distillation can overcome the limitations of L2L KD in token prediction recovery by reforming intermediate representations. Additionally, we found that applying token-wise logit scaling can significantly mitigate the risk of overfitting.
Drawing from our findings, we introduce a novel KD technique known as Token-Scaled Logit Distillation (TSLD), designed to enhance QAT for ternary quantization inference. We evaluate TSLD across a range of GLMs – originating from GPT-2 [2], OPT [4] and LLaMA [5] – of various sizes, including 7 billion models for the first time. The results show that TSLD achieves comparable, if not superior, performance in language modeling on ternary and 4-bit inference. When TSLD is applied to reasoning tasks, it surprisingly prevents overfitting to achieve task accuracy that is at least on par, if not better. These remarkable outcomes underline the potential of our proposed TSLD method in facilitating the deployment of ultra-low precision GLMs.
Fine-tuning for Generative Language Model. GLMs are renowned for their unparalleled text generation, comprehension, and reasoning capabilities [1-7]. Studies reveal that their performance can be enhanced through instruction fine-tuning methods like Prefix-Tuning [11] or using natural language instructions and examples [8]. Fascinatingly, instruction-tuned models, including smaller ones, can outperform larger counterparts in specific tasks [9, 8, 12]. However, the vast parameter count in these models, compared to popular models such as BERT [27], can restrict their practical utility. To mitigate this, we suggest investigating efficient, lightweight techniques for GLMs encompassing up to 7 billion parameters.
Quantization for Generative Language Model. Quantization, a method that minimizes the inference cost of large models by utilizing a limited number of bits to represent weights, has been recently applied to GLMs [13-17]. This process has considerably cut down GPU memory usage and execution time. Two main types of quantization exist, quantization-aware training (QAT) and post-training quantization (PTQ), which differ in their requirement for re-training or fine-tuning. Although QAT has been effectively used in Transformer encoder models [22, 25], its application to GLMs poses challenges [24], with observed performance declines when applied to decoder-only models like GLMs [25]. This paper assesses the imbalance of quantization errors based on attention traits and language model generation patterns. We demonstrate that QAT can be conducted without substantial performance loss across various tasks, even in models exceeding one billion parameters.
Knowledge Distillation for Language Model Compression. Knowledge distillation (KD) is a prevalent transfer learning framework that imparts knowledge from a larger "teacher" model to a smaller "student" model, and it is effectively used to curb accuracy decline in models compressed through quantization-aware training (QAT) [28, 19, 21-23]. In encoder models, KD trains the quantized model (student) using the full-precision model's (teacher's) intermediate representation as the objective in QAT. Despite its effectiveness, this method requires more memory due to the intermediate representation from the teacher model and has been less explored in decoder models such as GLM [25, 24]. This paper introduces a novel, less memory-intensive KD method applicable to models with up to 7 billion parameters. We offer a thorough analysis of the teacher's information transfer in the decoder model and suggest a QAT-based KD method that retains minimal performance degradation, even when applying ternary weights to various GLMs.
3 Background and Challenges
3.1 Transformer Layer
Generative language models [1] are built with Transformer layers [29]. A standard Transformer layer includes two main sub-modules: Multi-Head Attention (MHA) and Feed-Forward Network (FFN). Input to the  $l$ -th Transformer layer is  $\mathbf{X}_l \in \mathbb{R}^{n \times d}$  where  $n$  and  $d$  are the sequence length and hidden state size, respectively. Let  $N_H$  be the number of attention heads and  $d_h = d / N_H$ .  $\mathbf{W}_h^Q$ ,  $\mathbf{W}_h^K$ ,  $\mathbf{W}_h^V \in \mathbb{R}^{d \times d_h}$  are the weight parameters projecting  $\mathbf{X}_l$  into Query ( $\mathbf{Q}_h = \mathbf{X}_l\mathbf{W}_h^Q$ ), Key ( $\mathbf{K}_h = \mathbf{X}_l\mathbf{W}_h^K$ ), and Value ( $\mathbf{V}_h = \mathbf{X}_l\mathbf{W}_h^V$ ), respectively. The attention score ( $\mathbf{A}_h$ ) is computed with the dot product of the projected Query and Key ( $\mathbf{A}_h = \mathbf{Q}_h\mathbf{K}_h^\top$ ). The normalized version of this result is then passed through the softmax function and multiplied by the Value to get the output as head $h = \text{softmax}(\mathbf{A}_h / \sqrt{d_h})\mathbf{V}_h$ . Then, the output of the Multi-Head Attention (MHA) is defined as follows:
$$
\operatorname {M H A} \left(\mathbf {X} _ {l}\right) = \operatorname {C o n c a t} \left(\text {h e a d} _ {1}, \dots , \text {h e a d} _ {N _ {H}}\right) \mathbf {W} ^ {O}. \tag {1}
$$
FFN consists of two fully-connected layers with weight parameters  $\mathbf{W}^1$  and  $\mathbf{W}^2$ :
$$
\operatorname {F F N} \left(\mathbf {Y} _ {l}\right) = \operatorname {G e L U} \left(\mathbf {Y} _ {l} \mathbf {W} ^ {1} + b ^ {1}\right) \mathbf {W} ^ {2} + b ^ {2}. \tag {2}
$$
Therefore, the operations at the  $l$ -th Transformer layer can be defined as:
$$
\mathbf {Y} _ {l} = \mathbf {X} _ {l} + \operatorname {M H A} (\text {L a y e r N o r m} (\mathbf {X} _ {l})); \quad \mathbf {X} _ {l + 1} = \mathbf {Y} _ {l} + \operatorname {F F N} (\text {L a y e r N o r m} (\mathbf {Y} _ {l})). \tag {3}
$$
3.2 QAT with KD for Transformer Decoders
QAT emulates inference-time quantization during training to learn parameters robust to the quantization error. In particular, ternary quantization represents all the weight parameters  $(\mathbf{W}^Q,\mathbf{W}^K,\mathbf{W}^V,\mathbf{W}^O,\mathbf{W}^1,\mathbf{W}^2)$  into ternary values  $\mathbf{t}\in \{+1,0, - 1\} ^k$  along with a scale factor  $\alpha$  for sub-2bit inference at deployment. In this work, we follow the approach of TWN [30] that analytically estimates the optimal  $\alpha$  and  $\mathbf{t}$  to minimize  $\| \mathbf{w} - \alpha \mathbf{t}\| _2^2$ , where  $\mathbf{w} = \mathrm{vec}(\mathbf{W})$  and  $k$  is the number of elements of the weight parameters.
Due to aggressive bit-reduction, ternary quantization causes significant accuracy loss. KD can help compensate for accuracy degradation, where the original full-precision model works as a teacher to guide the training of the quantized model as a student. In case of Transformer models, prior works [19, 21, 22, 31] applied KD on every Transformer layer's output activation  $\mathbf{X}_l$  as well as attention scores  $\mathbf{A}_l$  with mean squared error (MSE) loss, denoted as  $L_{L2L}$ :
$$
L _ {L 2 L} = \sum_ {l = 1} ^ {L + 1} \operatorname {M S E} \left(\mathbf {X} _ {l} ^ {S}, \mathbf {X} _ {l} ^ {T}\right) + \sum_ {l = 1} ^ {L} \operatorname {M S E} \left(\mathbf {A} _ {l} ^ {S}, \mathbf {A} _ {l} ^ {T}\right), \tag {4}
$$
where superscripts  $S$  and  $T$  represent the student and teacher models, respectively.
The final output logits of the student  $(\mathbf{Z}^S)$  and the teacher  $(\mathbf{Z}^T)$  are used to compute the cross-entropy loss. Given that  $N$  is the total number of the tokens in input sequences, and  $V$  is the vocabulary size the language model can recognize and generate. Using the softmax function, we can convert each model's  $i_{th}$  token prediction logit output into probability distributions, which are then utilized for the loss for logit distillation  $(L_{logit})$ :
$$
\mathbf {P} _ {i} = \frac {e ^ {\mathbf {Z} _ {i , j}}}{\sum_ {j = 1} ^ {V} e ^ {\mathbf {Z} _ {i , j}}}, \quad L _ {\text {l o g i t}} = - \frac {1}{N} \sum_ {n = 1} ^ {N} \sum_ {i = 1} ^ {V} \mathbf {P} _ {n, i} ^ {T} \log \left(\mathbf {P} _ {n, i} ^ {S}\right). \tag {5}
$$
The overall loss for KD is generally computed as  $L_{KD} = L_{L2L} + L_{logit}$ , without GT Loss as noted in previous studies [19, 22, 31]. Yet, some methods utilize only  $L_{logit}$  [24]. Our study underscores the necessity of integrating  $L_{logit}$  and the GT Loss for an effective application of QAT in Transformer decoders.
Figure 1: (a) Illustration of attention mechanism in the encoder (top) and decoder models (bottom). (b) Left: performing NLU task [32] by encoder model. Right: performing language modeling task by decoder model with teacher-forcing (input token is independent of the previously generated token)
3.3 Quantization Challenges on GLMs
In this section, we compare the computations of Transformer encoders and decoders to deepen our understanding of the fresh challenges that surface within the realm of GLMs.
Cumulative Quantization Errors in Causal Attention. Causal attention, which integrates masking into self-attention to avoid referencing future tokens, is vital for causal language modeling tasks. To comprehend the quantization characteristics of GLMs, we contrast the computation procedures of the self-attention map. For a Transformer encoder, the quantization error reflected on the self-attention map due to the process of projecting Query, Key, and Value is evenly spread across tokens because of the token-parallel nature of computing inherent in Transformer encoders. However, the mask in the causal attention accumulates quantization errors from each token, creating uneven errors in the output computation.
Fig. 1(a) contrasts the attention mechanism of the encoder and decoder model under quantization. The encoder's self-attention is illustrated at the top of the figure, decomposed into self-attention and quantization error ( $Q_{err}$ ) components. The combined attention probabilities are utilized in a weighted sum with the Value, where tokens 1 and 3 (in a blue and red box respectively) are affected by an identical number of attention probabilities with quantization error. Conversely, the decoder's causal attention, shown at the bottom, uses only the attention probabilities of the current token and its preceding ones. For instance, the Value for token 1 in the bottom of Fig. 1(a) (in a blue box) uses only two attention probabilities affected by quantization error, while token 3 (in a red box) includes those from all preceding tokens. This illustration highlights that causal attention inherently leads to a disproportionate accumulation of quantization errors in the latter tokens. Thus, we need a decoder QAT strategy that addresses this imbalance within the causal attention module. In Section 4.1, we assess the limitations of current KD methods in managing cumulative quantization errors and introduce an enhanced logit distillation error compensation mechanism.
Necessity of Ground Truth Loss. In fine-tuning, encoder models, often used in Natural Language Understanding (NLU), and decoder models, common in Natural Language Generation (NLG), employ distinct mechanisms for receiving GT Loss, as shown in Fig.1(b). Encoder models for NLU tasks use a single special token to compute cross-entropy loss with a limited number of classes [27], as depicted on the left of Fig.1(b). On the other hand, decoder models in NLG tasks predict each subsequent token, transforming each token's representation into a logit vector with a class size equivalent to the vocabulary size, often exceeding 50k [1], shown on the right of Fig. 1(b). This process allows decoder models to obtain GT Loss for each input token, providing detailed token-level prediction information. Given these differences, there is a compelling need to consider the necessity of GT Loss in the decoder model's QAT in a token-wise manner. However, previous QAT [24] on the decoder models neglects the consideration of GT Loss due to a perceived degradation in performance when GT Loss is utilized. Accordingly, Section 4.2 offers an in-depth analysis of the interplay between KD and GT Loss during the QAT.
Figure 2: Comparison of weight quantization error on attention map through MSE loss in (a) encoder (BERT-base, RTE task) and (b) decoder (GPT-2, PTB task) model. (c) Left: min-max dynamic range per layer (Logit Distill vs L2L Distill). Middle: attention ranking ratio comparison (FP vs Logit Distill vs L2L Distill). Right: per layer token-wise logit distance and MSE distance
4.1 Logit Distillation for Cumulative Quantization Error
Motivation. The inherent nature of causal attention, where each token representation builds upon the representation of the previous tokens, presents previously unseen challenges when applying quantization to decoder models. For a clearer understanding of the decoder model, we conduct a comparative analysis with the encoder model to examine the impact of quantization error on the model. In Fig. 2 (a), the quantization error of the encoder self-attention map exhibits a widespread presence of errors due to the absence of masking in self-attention, and the per-token quantization errors along the layers also show irregular patterns depending on the token index. However, in Fig. 2 (b), the heat map of the decoder model reveals an increasing brightness of quantization errors as we move toward the later tokens. When examining the token index, the phenomenon of quantization errors accumulating toward the later tokens becomes even more pronounced. This previously unconsidered phenomenon of token quantization error accumulation in the decoder model is a crucial feature to consider in GLM QAT. Reflecting on this feature, we analyze the effectiveness of prior KD methods for language modeling and explore suitable KD approaches for the decoder model. Analysis on cumulative quantization error for a wider variety of GLMs can be found in Appendix A.3.
Comparison of KD Methods for Decoder QAT. Building on a deeper comprehension of the decoder model, we evaluate the efficiency of current KD methods for QAT in decoders and propose an enhanced KD approach informed by our decoder model analysis. We analyze how two different KD methods, Layer-to-Layer distillation (L2L KD) and logit distillation (Logit KD), tackle systematic outliers in QAT [15], using the min-max dynamic range per token and per channel of each layer's intermediate output. As shown in Fig. 2(c) left, both KD methods demonstrate distinct strategies in addressing the teacher model's systematic outliers. While L2L distillation guides the QAT process to mirror the outliers of the teacher model, Logit KD deviates from this pattern, generating new outliers not seen in the teacher model. These outliers consistently emerge in specific channel indices where the teacher model's outliers are present. Additionally, to compare the relative token attending order within each QAT model's self-attention map, we employ a ranking ratio comparison method [22]. This technique conveys the average relative importance of a single token within each attention map. As depicted in Fig. 2(c) middle, the L2L KD method closely mirrors the teacher model's ranking changes. However, the Logit KD method exhibits substantial variation in this ranking shift within a certain head range.
Figure 3: (a) Training/evaluation loss curve with different QAT-KD methods. (b) token-wise prediction statistics scatter plot (left: cross entropy loss and token confidence, right: cross entropy loss and token entropy). (c) Impact of TSLD: per-token scale and cross-entropy loss (left: per-token scale, Right: per-token cross-entropy loss). Analysis utilizes OPT-125m for PTB task language modeling. See Appendix A.2 for further analysis on other GLMs
Logit Distillation for Token-wise Prediction Recovery. We further analyze the QAT model's token logit distributions. Since token representations evolve along the layers to form the next token's probability [33], we assess each layer's logit distribution and the logit distance from the teacher model. As depicted in Fig. 2(c), L2L KD creates a token representation that closely mirrors the teacher model in both logit distribution and mean-squared error (MSE) distance during mid-layer stages but fails to match the final logit distribution. Conversely, Logit KD, despite diverging from the teacher model's logit distribution in the middle layers, accurately reproduces the final logit distribution. These observations highlight Logit KD's distinct mechanism for token-wise prediction recovery, managing cumulative quantization error in decoder models. In intermediate layers, Logit KD varies the attention values across channels as shown in Fig.2(c), leading to a diverging token representation from the FP model, with this middle stage adjustment acting to counteract accumulated quantization errors in later tokens. Consequently, Logit KD aligns the final logit distribution for each token, crucial for the accuracy of causal language modeling. Therefore, Logit KD, aligning with the characteristics of the decoder model, stands out as a natural choice for QAT. The subsequent section will delve into previously unexamined issues encountered by Logit KD in decoder QAT.
4.2 Token-Scaled Logit Distillation for Avoiding Overfitting with GT Loss
Motivation. This section tackles the overfitting problem arising from the combination of Logit KD and GT Loss during QAT. We also investigate the probabilistic behavior displayed by the decoder model in language modeling tasks. The study by [24] highlights instances where employing GT Loss and Logit KD adversely impacts the performance of decoder QAT. To understand this issue better, we conduct tests using Logit KD both independently and combine with ground truth loss. As depicted in Fig. 3 (a), overfitting is observed in the QAT when both ground truth loss and Logit KD are applied.
Understanding Causes of Overfitting. To better understand the causes of overfitting, we analyze the logit output for each token that the teacher model generates during language modeling. From the logit information by the teacher model, we derive the probability distribution  $(\mathbf{P}_i^T = \mathrm{softmax}(\mathbf{Z}_i^T))$  for  $i_{th}$  token prediction. Based on this distribution, we further calculate cross-entropy  $(-y_{n,i}\log (\mathbf{P}_{n,i}^{T}))$  confidence score  $(\max (\mathbf{P}_i^T))$  and entropy  $(- \sum_i\mathbf{P}_i^T\log (\mathbf{P}_i^T))$  for each token prediction. These
metrics unveil a confidence disparity in language modeling—a trend uniformly observed across decoder models of varying scales. The cross-entropy loss of token prediction logits plotted against the probability confidence score, as illustrated in Fig. 3 (b), distinctly demarcates the Low Confidence Region (blue box) with low probability confidence and high cross-entropy loss from the High Confidence Region (yellow box) with high probability confidence and low cross-entropy loss. This observation implies a potential overlap between high-confidence Logit KD and the role of GT Loss cross-entropy, suggesting that redundant information from high-confidence Logit KD might mirror the effects of ground truth loss, thereby contributing to the overfitting observed.
Token-Scaled Logit Distillation (TSLD). Based on investigations into the probabilistic relation of token predictions and overfitting in QAT, we propose an adaptive KD method that adjusts Logit KD based on token confidence. This approach utilizes the phenomenon of confidence disparity in token predictions from the teacher model. Our method, called Token-Scaled Logit Distillation (TSLD), de-emphasizes Logit KD for high-confidence tokens to prevent overfitting in QAT while emphasizing Logit KD for low-confidence tokens possessing with a high entropy probability distributions. Specifically, low-scaled Logit KD (high confidence, low entropy) effectively reduces the overlap with GT Loss, leading to an improvement in overfitting. On the other hand, high-scaled logit KD (low confidence, high entropy) emphasizes the distillation of more informative token prediction distribution from the teacher model, which has rich soft label information.