Compression of Generative Pre-trained Language Models via Quantization

Introduction

Transformer-based generative pre-trained language models (PLMs) show strong abilities of multi-task and few-shot learning, and achieve remarkable performances on various tasks (Radford and Narasimhan 2018; Brown et al. 2020; Lewis et al. 2020; Raffel et al. 2020; Chen et al. 2021). However, they are usually expensive in terms of both computation and memory due to a large number of parameters, and the token-by-token generation process. Many methods have been proposed to compress PLMs, but mostly focus on understanding tasks like sentence classification with BERT (Lan et al. 2019; Sun et al. 2020b; Jiao et al. 2020; Shen et al. 2020; Hou et al. 2020). Recent works try to compress GPT-2 using tensor decomposition (Edalati et al. 2021), and knowledge distillation (Song et al. 2020), but the compression ratio achieved is much smaller than that of BERT. Yet the underlying difficulty remains unclear.

In this paper, we firstly explore compressing generative PLMs by quantizing the parameters from full-precision to lower bits. We find that directly applying previous quantization methods designed for BERT or computer vision tasks to generative PLMs lead to poor performance. Figure 1 shows that the performance drops sharply as the weight bit-width decreases. To investigate the difficulty of quantizing generative PLMs, we find that the learned embeddings tend to be homogeneous and hard to distinguish due to the reduced capacity caused by quantization, while the weight distributions also vary significantly across different modules and different Transformer layers. These problems are further magnified due to the nature of sequential left-to-right prediction of generative PLMs, as the quantization error will accumulate across time.

To alleviate the above problems, we propose a token-level contrastive distillation to contrast on tokens and make the word embedding distinguishable. Besides, we propose a module-wise dynamic scaling for the quantizer to better adapt to different modules. Empirical results on language modeling, next utterance prediction and summarization show that compared to the full-precision baseline, our quantized GPT and BART (abbreviated as QuantGPT and QuantBART) achieve comparable performance for 8/4-bit weight, and have only a slight drop for 2-bit weight, while being over 13$\times$ smaller. QuantGPT also clearly outperforms previous GPT compression methods on language modeling.

To summarize, our main contributions are: 1) We find that generative PLMs are hard to quantize due to homogeneous word embedding and varied weight distribution. 2) We then propose the token-level contrastive distillation and module-wise dynamic scaling, to make the word embedding more distinguishable and make quantizers adapt to different modules, respectively. 3) Empirical results on various tasks show the efficacy of our method.

Difficulty of Qunatizing Generative Pre-trained Language Models

In this section, we show that it is challenging to train a low-bit generative pre-trained model with conventional quantization approaches directly. Before diving into details, we first review the necessary backgrounds of quantization.

In this paper, we apply the quantization-aware training Courbariaux et al. 2015 to generative PLMs. Specifically, denote the vectorized full-precision weight as $\mathbf{w}$, each forward propagation first clips the weight by a positive clipping factor $\alpha$, and then quantizes the clipped weight to $b$-bit as

A good clipping factor is expected to take the majority of full-precision weight into account via clipping, i.e., quantizing the range where data are densely distributed to reduce quantization error. To solve this problem, PACT (Choi et al. 2018) learns a parameterized clipping factor and achieves better results than setting a fixed clipping factor. Instead of learning the clipping factor, LSQ (Esser et al. 2020) learns the step size $\alpha/n$, but requires a careful initialization and gradient update.

In practice, following previous works on BERT quantization (Zhang et al. 2020; Bai et al. 2021), we use layer-wise quantization (i.e., one clipping factor for elements in each weight matrix) for all weight matrices in the Transformer layers and row-wise quantization (i.e., one clipping factor for each word embedding) for the embedding layer. We use asymmetric uniform quantization for activations after self-attention and GeLU function whose elements are mostly positive, and symmetric uniform quantization for other activations. We do not quantize layer-normalization layers, skip connections, biases due to small computational overhead.

2 Difficulty Analysis

We compare the following representative quantization methods including (i) LAQ (Zhang et al. 2020) for BERT; (ii) PACT (Choi et al. 2018) and LSQ (Esser et al. 2020)) for computer vision tasks, to generative pre-trained model, GPT-2. Figure 1 shows the performance under different weight bit-widths, and the performance drops sharply as the bit-width decreases, especially for PACT and LSQ. In the following, we study the potential reasons behind the difficulty of quantizing generative PLMs, by empirically investigating the properties of the word embedding and model parameters.

We first study the difficulty from the learned word embeddings of different models. In Figure 2, we visually compare the distributions of the word embeddings of the full-precision and quantized models under the same scale. As can be seen, the word embeddings of the full-precision model are scattered distinguishable, while those in previous quantization methods PACT, LSQ and LAQ learn homogeneous word embeddings which are clustered and less distinguishable, especially for PACT and LSQ. We speculate this is caused by the sequential computation nature of GPT. Specifically, unlike BERT which computes the representation of all tokens in parallel, GPT computes each token in left-to-right order, and the quantization error incurred in the previous tokens will pass on to future tokens, making the learning signal noisier over time, and finally less informative word embeddings.

A direct consequence of the homogeneous word embedding can be reflected in Figure 3. By comparing Figure 2 and Figure 3, we can find that the higher degree of homogeneity in the word embedding of a quantized model, the fewer dependencies among different tokens are kept.

As will be discussed in Section 3.1, we propose a token-level contrastive learning to alleviate this problem. Compared with PACT, LSQ and LAQ, our method not only aligns the token representations between the quantized and full-precision networks (i.e., diagonal boxes), but also captures the dependencies among different tokens (non-diagonal boxes). More visualizations are available in Appendix C.3. The non-distinguishable word embeddings and poor ability to capture contextualized dependencies also make methods like PACT and LSQ more likely to generate incorrect tokens, e.g. illogical and repeated text ( Section 4.4).

Besides the learned word embeddings, we also investigate the distribution of the weights in the full-precision model. Figure 4 shows that the weight distributions of a 12-layer full-precision GPT-2 are highly skewed with outliers. This causes difficulty in estimating the clipping factor $\alpha$ of the quantizer by heuristic methods, or even by PACT which learns the $\alpha$ through gradient descent. Specifically, in PACT, the approximated gradient of $\alpha$ only relies on the weights whose absolute values are larger than $\alpha$. This solution ignores the effect of weights within $[-\alpha,\alpha]$ and depends heavily on the initialization of $\alpha$. Figure 4 shows that an improper initialization together with the inaccurate gradient estimation of the clipping factor often make the learned $\alpha$ of PACT too large, and can not provide fine resolution to the majority of weights within the clipping range. The quantization error accumulated over time makes this problem more severe. In this work, we re-parameterize the clipping factor to make the quantizer adaptive to each module in the Transformer layers, and consider both weights outside and inside the clipping range when estimating the gradient of the clipping factor.

As will be discussed in Section 3.2, we propose a module-wise dynamic scaling to reduce the clipping factor’s sensitivity to initialization, and an improved gradient estimation that also considers the weights within $[-\alpha,\alpha]$. Figure 4 shows that the clipping factor learned by our method gives finer resolutions to the majority of the weights.

Proposed Method

Based on the observations in Section 2, we propose a quantization method which utilizes token-level contrastive distillation to make the word embedding distinguishable (Section 3.1) and a module-wise dynamic scaling adjustment to learn better clipping factors (Section 3.2).

The proposed token-level contrastive distillation contrast among tokens instead of sequences sequence, to learn distinguishable representations for each token. Inspired by Baevski et al. 2020, which uses in-utterance representation at different positions of the same utterance as negatives for speech feature learning, for each token of the quantized network, we use the representation of the same token from the full-precision teacher network as its positive, while representations of other tokens in the same sequence as negatives (Figure 5). Inspired by He et al. 2020 which uses a momentum encoder for more consistent representation, we build a memory bank to store momentum token representations. When computing the contrastive distillation loss, we load the representations of negative samples from the memory bank with cheap indexing operations.

where $\lambda$ is a trade-off factor set as 0.1 by default.

The proposed token-level contrastive distillation outperforms the sequence-level counterpart (as will be shown empirically in Section 5.1.1). We conjecture this is because (i) token-level contrast alleviates the problem of homogeneous word embedding (Figure 2) in the low-bit quantization; and (ii) similar to speech, the order of natural language is also sequential instead of spatial like images; and (iii) the self-attention mechanism allows other tokens to learn representations contextualized on the studied token, and these in-sequence negatives are harder than those from in-batch sequences, allowing more efficient representation learning.

2 Module-dependent Dynamic Scaling

Based on the observation of varied weight distribution in Section 2, we propose a simple-yet-effective dynamic scaling according to the statistics of each module weight. Specifically, instead of directly learning the original clipping factor $\alpha$ as PACT, we turn to learn a new scaling factor $\gamma$, which is multiplied with the average weight magnitude $\frac{\|\mathbf{w}\|_{1}}{n}$ to get clipping factor $\alpha$:

Experiments

In this section, we evaluate the efficacy of our proposed quantization method on three kinds of generative tasks on two kinds of generative pre-training models. Specifically, we perform the proposed quantization approach on language modeling and next utterance prediction tasks on GPT-2 (Radford and Narasimhan 2018), and abstractive summarization using BART (Lewis et al. 2020), and call the resultant models QuantGPT and QuantBART. The token-level contrastive distillation is performed on the hidden states of the last layer of GPT-2 or the BART decoder. More details about the datasets and model architectures can be found in Appendix B.1 and B.2.

For each downstream task with our proposed method, we first fine-tune a full-precision network using the pre-trained checkpoint from huggingface http://huggingface.co/models for both GPT-2 and BART. Then we use this fine-tuned network as the full-precision teacher network and to initialize the quantized student network. We train each task with 8 V100 GPUs based on the Pytorch framework. The detailed hyper-parameters for each task are available in Appendix B.3.

Since there are very few attempts to compress generative PLMs, we self-implement three baseline quantization methods PACT (Choi et al. 2018), LSQ (Esser et al. 2020) and LAQ (Hou and Kwok 2018) for comparison. Details about these methods are in Appendix B.4.

2 Language Modeling

The task of language modeling is to predict the probability distribution over a sequence of words. For language modeling, we experiment on WikiText2 (Merity et al. 2016), Penn Treebank (PTB) (Mikolov and Zweig 2012) and WikiText103 (Merity et al. 2016). We use perplexity (PPL) to evaluate the performance for language modeling.

From Table 1, the performance of the proposed method with 8-bit weight is comparable to the full-precision counterpart on PTB and WikiText103, while drops slightly on WikiText2. A slightly more severe performance drop is observed as the bit-width decreases from 8 to 4, with a drop of around 1 PPL point on WikiText2 and WikiText103, and less than 0.1 PPL point on PTB. When the bit-width of weight further goes down to 2, our method has an average of 2 PPL points drop, but achieves 14.4$\times$ model size reduction.

From Table 1, our method outperforms PACT, LSQ and LAQ for all bit-widths and tasks. As the bit-width decreases from 8 to 4, the PPL of LSQ greatly increases, with the average PPL of LSQ increasing by over 5 times. As the bit-width further decreases to 2, both LSQ and PACT fail on all datasets, despite their good performance on understanding tasks on BERT (Bai et al. 2021). We conjecture it is because though both PACT and LSQ have learnable parameters, the accumulated quantization error of generative PLMs makes the updates of these parameters by gradient descent less stable. On the other hand, the proposed module-wise dynamic scaling alleviates the problem.

In Table 2, we compare our quantization method against recent GPT-2 compression methods, including tensor decomposition method KnGPT2 (Edalati et al. 2021), as well as distillation methods DistilGPT2 and LightPAFF (Song et al. 2020). From the comparison, our method outperforms the others in terms of model size and performance, even when weights are compressed to only 2 bits.

3 Next Utterance Prediction

The task of next utterance prediction predicts the next utterance given the dialogue context. It tests the language understanding ability of generative models. For this task, we use a large-scale dialogue dataset, Persona-Chat (Zhang et al. 2018).

From Table 1, all quantization methods incur a clear performance drop compared to the full-precision baseline, even in the 8-bit setting. As the quantization becomes more aggressive, i.e., the bit-width gets smaller, the performance of PACT and LAQ decrease more significantly than ours. In particular, LSQ diverges for 2-bit weight and its accuracy is only 5%, which is no better than a random guess as there are 20 classes.

4 Abstractive Summarization

Abstractive summarization aims at generating a terse summary that captures the main ideas of the source article. We experiment on XSum (Narayan et al. 2018), whose ground-truth summarizations are highly abstractive and are challenging for many extractive strategies. ROUGE 1, 2, L are used to evaluate the performance of this task.

Table 3 shows the results of the abstractive summarization. As can be seen, our method constantly outperforms other methods again with a clear margin. Example generated summarizations of different methods in Appendix C.2 show that the summaries generated by QuantBART are logical and terse, while those from PACT have repeated texts.

Discussion

From Table 4, “fp+quan.” and “quan. only” performs worse than QuantGPT, which uses full-precision representations of other tokens as negative samples. This indicates that noisy representations of tokens from the not-fully-trained quantized network may not be sufficient. “global” performs even worse, which we conjecture is because, for one token, negative tokens chosen from the same sequence are contextually related to it and more informative than random tokens. “in-batch” performs worse than all token-level variants, which may be because generative tasks make predictions in a token-wise manner and rely heavily in finer-grained token-wise representations. Interestingly, contrary to in-batch negative sampling in computer vision (Chen et al. 2020), we find that reducing the number of negative samples by reducing the batch size from 32 to 16 slightly improves performance.

1.2 Number of Negative Samples

In Figure 8, we plot the PPL of 2-bit QuantGPT on the PTB dataset, with varying number of negative samples. We plot the mean results with standard deviations from 5 independent runs. As can be seen, the performance improves and converges gradually as the number of negative samples increases. Figure 8 also shows that using the moving-average representations (${\bf q}^{s}_{t_{i}}$ in Eq. (3)) of negative samples in the memory bank has better performance than using the immediate representations (${\bf h}^{s}_{t_{i}}$ in Eq. (3)), because of a smoother and more consistent representation of tokens.

1.3 Training Cost of the Contrastive Loss

In Table 5, we report the training speed and memory consumption of training the GPT-2 model on the PTB dataset with and without the proposed token-level contrastive loss. Batch size is set to 4 per device, which can be increased by using GPUs with larger memory or reducing the sequence length of samples. As can be seen, with the proposed token-level contrastive loss, the performance clearly improves with only slightly slower training speed and more memory consumption.

1.4 Representations for the Contrastive Loss

In Table 6, we compare the different representations to perform the contrastive loss. The “decoder-last”( resp. “decoder-first”) denotes performing the proposed token-level contrastive loss on the hidden states from the last decoder layer (resp. first decoder layer) followed by a linear transformation.

2 Ablation on Dynamic Scaling

Figure 8 shows the learned scaling $\gamma$ of different modules in the 2-bit GPT-2 model. As can be seen, the scalings of different modules vary a lot, verifying the need for module-wise dynamic scaling.

Related Work

Some early explorations compress the generative pre-trained language models. KnGPT2 (Edalati et al. 2021) applies the Kronecker decomposition to compress the GPT. DistilGPT2 https://transformer.huggingface.co/model/distil-gpt2 distills a 12-layer GPT-2 to a 6-layer one, which is twice as fast during inference. LightPAFF (Song et al. 2020) proposes a distillation approach that the training loss is a combination of a maximum likelihood loss of the student model, and the KL divergence between the output of teacher and student models. SpAtten Wang et al. 2021 proposes a sparse model with algorithm and architecture co-design, which removes uninformative tokens and attention heads. Compared with these methods, we not only study the difficulties of compression from the properties of generative tasks, but also study both decoder and encoder-decoder generative models.

Quantization compresses a model by representing the 32-bit floating-point parameter with a low-bit representation, and has been widely used in various domains as it does not require designing a new model architecture. There have been many attempts to quantize task-specific BERT models (Zafrir et al. 2019; Shen et al. 2020; Zadeh et al. 2020) with only negligible performance drop on natural language understanding tasks. Recent works (Zhang et al. 2020; Bai et al. 2021) even push the weight bit-width down to as low as 1-bit. Despite the success of these approaches for BERT models, attempts to quantize generative PLMs are scarce, and the underlying difficulty remains unclear.

Contrastive learning aims at pushing the representations of similar samples together while pulling those of dissimilar ones apart. and is widely used for large-scale self-supervised learning in various domains (Chen et al. 2020; Sun et al. 2020a; Baevski et al. 2020; Huang et al. 2022), and multi-modal learning (Radford et al. 2021; Jia et al. 2021). SimCLR (Chen et al. 2020) directly uses other in-batch samples as negatives, and sufficient large batch size is required to work well. MoCo (He et al. 2020) maintains a large number of negative samples in a queue and uses a moving average key encoder to improve consistency. Contrastive learning without negative samples is also proposed in BYOL (Grill et al. 2020) and SimSiam (Chen and He 2021). Contrastive representation distillation (Tian et al. 2019) distills the knowledge from the teacher network to the student network by maximizing the mutual information between them.

The closest work with our token-level contrastive distillation is Wav2vec 2.0 (Baevski et al. 2020), which use in-utterance representations at different positions as negatives in speech learning. Besides the difference in the modality and tasks, our method also differs from theirs in (1) Model: We quantize the model parameters and activations while they do not; (2) Representation: For each sample, we use the output of the full-precision and the quantized networks as its two views, while they use the quantized and the contextualized representation. (3) Loss: We calculate loss over all tokens in an auto-regressive manner, while they only calculate over the masked tokens non-autoregressively.

Conclusion

This paper studies low-bit quantization of generative PLMs. We find that the difficulty of quantizing generative PLMs lies in homogeneous word embedding and varied distribution of weights. To alleviate the two problems, we propose token-level contrastive learning to learn more distinguishable token embeddings, as well as a module-dependent dynamic scaling for more accurate quantization. Extensive experiments on language modeling, next utterance prediction and abstractive summarization demonstrate the efficacy of our proposed method. We hope our work sheds a light on the compression of generative PLMs in future exploration.

Limitations

Acknowledgements

This work is supported in part by the General Research Fund (GRF) project 17206020, and in part by ACCESS, AI Chip Center for Emerging Smart Systems, Hong Kong SAR.

References

Appendix A Derivation of Gradient of Dynamic Scaling

In this section, we provide the derivation of the gradient of the proposed dynamic scaling $\gamma$. The quantization in the forward process can be written as

We use straight through estimator (STE) to estimate the gradient of uniform quantizer $Q(\cdot)$, i.e., $\forall i,\frac{\partial Q(u_{i})}{\partial u_{i}}=1$. Thus the gradient $\frac{\partial Q(u_{i})}{\partial\alpha}$ can be written as:

which considers both the weight inside and outside the clipping value, and is proportional to the weight magnitude $\frac{\|\mathbf{w}\|_{1}}{n}$.

Appendix B More Experimental Settings

The train/val/test splits for different datasets are shown on Table 8.

B.2 Model Architectures

The vocabulary size of GPT-2 is 50527. We use GPT-2-small with 12 decoder layers and hidden state dimension of 768, for experiments in Sections 2.2, 1 and 5. GeLU (Hendrycks and Gimpel 2016) is used as the activation function. In the experiments of Appendix C.1, we adopt GPT-2-base with 24 decoder layers and hidden state dimension of 1024, to evaluate the quantization ability on larger models.

The vocabulary size of BART is 50265. We use BART-base with 6 encoder layers, 6 decoder layers and hidden state dimension as 768 for experiments in Section 1. In the experiments of Appendix C.1, we adopt BART-large with 12 encoder layers, 12 decoder layers and hidden state dimension 1024, to evaluate the quantization ability on larger models.

B.3 Hyperparameters

The sequence length is 512. The learning rate is initialized to 0.0005 (resp. 0.001) for the GPT-2 backbone parameters (resp. clipping factor $\gamma$) and then linearly decays to 0. The number of negative samples in each sequence is 64 for the PTB dataset, and 32 for the WikiText2 and WikiText103. The temperature $\tau$ and momentum coefficient $m$ is 0.1 and 0.5 respectively. We train with the AdamW optimizer Loshchilov and Hutter 2017 with batch size 32. The training epochs for WikiText2, PTB and WikiText103 are set as 80, 120, 8, respectively.

The sequence length is 512. The learning rate is initialized to 0.0005 (resp. 0.001) for the GPT-2 backbone parameters (resp. clipping factor $\gamma$) and then linearly decays to 0. The number of negative samples in each sequence is 32. The temperature $\tau$ and momentum coefficient $m$ is 0.1 and 0.5, respectively. We train with the AdamW optimizer with batch size 16, for a total of 2 epochs.

We set the length of the source sequence (articles) as 512, and pad the target sequence (summaries) to maximum length. We use beam search to generate summaries, with beam size 6 and length penalty 1. The learning rate is initialized to 0.0002 (resp. 0.001) for the BART backbone parameters (resp. clipping factor $\gamma$) and then linearly decays to 0. The number of negative samples is 32. The temperature $\tau$ and momentum coefficient $m$ is 0.1 and 0.5, respectively. We train with the AdamW optimizer with batch size 128, for a total of 8 epochs.

B.4 Description of the Compared Methods

PACT (Choi et al. 2018) learns a learnable clipping factor for each module by gradient descent. To make the quantization more accurate, we adopt a flexible variant of the original PACT, with different positive and negative clipping factors $[-\alpha_{neg},\alpha_{pos}]$, where both $\alpha_{neg}$ and $\alpha_{pos}$ are initialized as 2.5.

LSQ (Esser et al. 2020) learns the step-size of quantizer for each module by gradient descent. We use the recommended initialization strategy of the step size as (Esser et al. 2020).

LAQ (Hou et al. 2017; Hou and Kwok 2018) is a loss-aware quantization method that views quantization as an optimization problem and solve it via proximal Newton algorithm. We use the approximate solver in (Hou and Kwok 2018) to compute the quantized weights before each forward propagation.

For the self-implemented methods PACT, LSQ and LAQ, we adopt the commonly-used distillation loss adopted in (Hinton et al. 2015; Jiao et al. 2020). Note that these methods are only used for weights and embeddings, while the activations of these methods follow the same setting as our proposed method in Section 2.1. We also tried using the original language modeling loss w.r.t. the ground-truth labels, and distillation loss over the attention as (Jiao et al. 2020). However, these two losses worsens the performance on all three methods.

B.5 Frameworks of Double-head GPT-2 and BART

Since we adopt double-head GPT-2 and BART for next utterance prediction and abstractive summarization, the frameworks for these tasks are slightly modified from that on language modeling due to the difference of tasks. In Figure 9 and 10, we illustrate the framework for double-head GPT-2 and BART, respectively. In the double-head GPT-2, we also quantize the final linear layer in the output head.

Appendix C More Experimental Results

In Table 9, we experiment with GPT-base and BART-large, which both have 24 Transformer layers. For all bit-widths, the training of our method converges successfully without gradient exploding/vanishing problems. QuantGPT outperforms PACT by a large margin in all tasks, especially for 2-bit weight. Our quantization method on larger models also has better performance than that on 12-layer GPT-2 and 12-layer BART in Section 1.

C.2 Examples of Summarizations

In Table 10, we provide the example summarizations on the XSum dataset. By comparing the articles, references and generations, the generated summaries by our quantized model are more logical and terse than PACT, LSQ and LAQ, which face problems of homogeneous word embeddings to some extent as discussed in Section 2.

C.3 More Visualizations for the Token Representations

In Figure 11, we provide the visualizations of token representations on more samples. The observations are similar to those in Section 2.