Token-Scaled Logit Distillation for
Ternary Weight Generative Language Models

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}^{Q}_{h},\mathbf{W}^{K}_{h},\mathbf{W}^{V}_{h}\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}^{Q}_{h}$), Key ($\mathbf{K}_{h}=\mathbf{X}_{l}\mathbf{W}^{K}_{h}$), and Value ($\mathbf{V}_{h}=\mathbf{X}_{l}\mathbf{W}^{V}_{h}$), respectively. The attention score ($\textbf{A}_{h}$) is computed with the dot product of the projected Query and Key ($\text{{A}}_{h}=\textbf{Q}_{h}\textbf{K}^{\top}_{h}$). The normalized version of this result is then passed through the softmax function and multiplied by the Value to get the output as $\text{\text{head}}_{h}=\text{softmax}(\text{{A}}_{h}~{}/{\sqrt{d_{h}}})\textbf{V}_{h}$. Then, the output of the Multi-Head Attention (MHA) is defined as follows:

$\displaystyle\text{MHA}(\mathbf{X}_{l})=\text{Concat}(\text{\text{head}}_{1},...,\text{\text{head}}_{N_{H}})\mathbf{W}^{O}.$

(1)

FFN consists of two fully-connected layers with weight parameters $\mathbf{W}^{1}$ and $\mathbf{W}^{2}$:

$\displaystyle\text{FFN}(\mathbf{Y}_{l})=\text{GeLU}(\mathbf{Y}_{l}\mathbf{W}^{1}+b^{1})\mathbf{W}^{2}+b^{2}.$

(2)

Therefore, the operations at the $l$-th Transformer layer can be defined as:

$\displaystyle\mathbf{Y}_{l}=\mathbf{X}_{l}+\text{MHA}(\text{LayerNorm}(\mathbf{X}_{l}))\text{;}\quad\mathbf{X}_{l+1}=\mathbf{Y}_{l}+\text{FFN}(\text{LayerNorm}(\mathbf{Y}_{l})).$

(3)

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 t to minimize $\|\mathbf{w}-\alpha\mathbf{t}\|^{2}_{2},$ where $\mathbf{w}=\text{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 $\textbf{X}_{l}$ as well as attention scores $\textbf{A}_{l}$ with mean squared error (MSE) loss, denoted as $L_{L2L}$:

$\displaystyle L_{L2L}=\sum_{l=1}^{L+1}\text{MSE}(\textbf{X}^{S}_{l},\textbf{X}^{T}_{l})+\sum_{l=1}^{L}\text{MSE}(\mathbf{A}^{S}_{l},\textbf{A}^{T}_{l}),$

(4)

where superscripts $S$ and $T$ represent the student and teacher models, respectively.

The final output logits of the student ($\textbf{Z}^{S}$) and the teacher ($\textbf{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}$):

$\displaystyle\textbf{P}_{i}=\frac{e^{\textbf{Z}_{i,j}}}{\sum_{j=1}^{V}e^{\textbf{Z}_{i,j}}},\quad L_{logit}=-\frac{1}{N}\sum_{n=1}^{N}\sum_{i=1}^{V}\textbf{P}^{T}_{n,i}\log(\textbf{P}^{S}_{n,i}).$

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

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.

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.

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.

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 ($\textbf{P}_{i}^{T}=\text{{softmax}}(\text{{{Z}}}_{i}^{T})$) for $i_{th}$ token prediction. Based on this distribution, we further calculate cross-entropy ($-y_{n,i}\log(\textbf{P}_{n,i}^{T})$), confidence score ($\max(\textbf{P}_{i}^{T})$) and entropy ($-\sum_{i}\textbf{P}^{T}_{i}\log(\textbf{P}^{T}_{i})$) 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.

$$CE^{T}_{n}=-\sum_{i=1}^{V}y_{n,i}\log(\textbf{P}_{n,i}^{T}),\quad\text{scale}_{n}=\frac{e^{\text{CE}^{T}_{n}/\tau}}{\sum_{k=1}^{N}e^{\text{CE}^{T}_{k}/\tau}}$$

(6)

$\displaystyle L_{TSLD}=\sum_{n=1}^{N}\left(\text{scale}_{n}\times-\sum_{i=1}^{V}\textbf{P}_{n,i}^{T}\log(\textbf{P}_{n,i}^{S})\right)$

(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 ($\text{scale}_{n}$) based on the cross entropy loss of the teacher model ($CE^{T}_{n}$) 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 ($\text{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]), common-sense QA tasks (PIQA [35], OpenbookQA [36], ARC_easy [37], ARC_challenge [37] and arithmetic 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 pre-processing 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.

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

We assess the effectiveness of our proposed method in commonsense QA (PIQA, OpenbookQA, ARC_easy, 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.

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.

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.

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.