Zero-shot Sharpness-Aware Quantization for Pre-trained Language Models (2023)

Abstract: Quantization is a promising approach for reducing memory overhead and accelerating inference, especially in large pre-trained language model (PLM) scenarios. While having no access to original training data due to security and privacy concerns has emerged the demand for zero-shot quantization. Most of the cutting-edge zero-shot quantization methods primarily 1) apply to computer vision tasks, and 2) neglect of overfitting problem in the generative adversarial learning process, leading to sub-optimal performance. Motivated by this, we propose a novel zero-shot sharpness-aware quantization (ZSAQ) framework for the zero-shot quantization of various PLMs. The key algorithm in solving ZSAQ is the SAM-SGA optimization, which aims to improve the quantization accuracy and model generalization via optimizing a minimax problem. We theoretically prove the convergence rate for the minimax optimization problem and this result can be applied to other nonconvex-PL minimax optimization frameworks. Extensive experiments on 11 tasks demonstrate that our method brings consistent and significant performance gains on both discriminative and generative PLMs, i.e., up to +6.98 average score. Furthermore, we empirically validate that our method can effectively improve the model generalization.

## introduction pre-trained language models (plms), such as bert @xcite and gpt-3 @xcite , have achieved great success in a variety of nlp tasks @xcite @xcite . however, with the scaling of model size, the inference of larger plms, e.g., opt @xcite and llama @xcite , becomes more computationally expensive and memory-intensive (shen et al., * miaoxi zhu and qihuang zhong contribute equally to this work. to achieve this goal, various model compression methods have been developed, including knowledge distillation @xcite , pruning @xcite and etc. among these methods, quantization has attracted great attention (especially in the large language model scenarios) @xcite , owing to its impressive ability to reduce memory footprint and accelerate inference while maintaining the network structure @xcite . however, quantization-aware training (qat) usually requires retraining using the original training data to mitigate performance degradation. while in some application scenarios, access to the original training data is oftentimes not available due to the protection of privacy and consideration of data security. in response to this problem, post-training quantization (ptq) @xcite is proposed to address the training-data access, as it generally requires no retraining. but most of those methods would lead to an accuracy drop in lowprecision quantization since they are training-free and primarily rely on the analysis of weight distributions. more recently, zero-shot quantization (zsq) @xcite @xcite shows promising results in the computer vision community. specifically, zsq performs the quantization process with the synthetic data generated from a generator. to alleviate the side effect of inaccurate synthetic data, @xcite further improve the zsq by adaptively tuning the generator with the supervision of adversarial learning. however, it is non-trivial to adopt such an adversarial-based zsq in the nlp field, as ❶ backpropagation on discrete words is not reasonable, i.e., it seems "unlikely" to pass the gradient through the text to the generator; ❷ minimizing the difference between the teacher and (quantized) stu-dent models would be prone to over-fitting problem, leading to poor model generalization. to address the above issues, we propose a novel zero-shot sharpness-aware quantization (namely zsaq) framework to improve the performance and generalization of the quantized model. firstly, regarding the problem ❶, we design a simple feature adaptation module to convert the output token representations of the generator into the target quantized model's representation space, thus avoiding the gradient propagation on the discrete words. secondly, for the problem ❷, we are inspired by the sharpness-aware minimization (sam) @xcite and propose an alternating sam-sga algorithm to boost the zsaq and improve model generalization. specifically, on the one hand, sam-sga uses sam to robustly minimize the divergence of output distributions between teacher and quantized student models. on the other hand, it uses stochastic gradient ascent (sga) to encourage the generator model to maximize the divergence at each iteration. by optimizing such a minimax problem, we can not only train the well-performed with generative adversarial learning but also improve the generalization of quantized models. theoretically, we provide rigorous analysis for the sam-sga algorithm solving the minimax optimization (zsaq). we show that sam-sga achieves o(1/ √ t ) convergence rate. furthermore, our theoretical results are not limited to the single case but are applicable to a wide range of minimax optimization problems in both nonconvexstrongly-concave and nonconvex-pl conditions. empirically, we adopt our zsaq framework to quantize both the discriminative and generative plms, and evaluate 8 language understanding tasks and 3 language modeling tasks. extensive experiments demonstrate the effectiveness and versatility of our approach. more encouragingly, our zsaq brings +6.98 average performance gains in the lowprecision quantization, compared to the baselines. extensive analyses show that our framework has the potential to expand to more large language models and prove that sam-sga indeed brings better model generalization. in summary, our contributions are three-fold: (1) we propose a novel method, zero-shot sharpnessaware quantization (zsaq), which realizes zeroshot quantization without much accuracy loss. (2) we provide a theoretical convergence guarantee for the sam-sga algorithm which aims to solve the minimax optimization problem (zsaq). (3) extensive experiments show that our sam-sga can bring consistent and significant performance gains, up to +6.98 average score, on both discriminative and generative plms. ## related works compression method for language models. compression is efficient in reducing computation cost, memory overhead, and energy consumption as well as shortens inference time for large neural network models, which is in great need for pre-trained language models to be deployed in the mobile or resource-constrained device. quantization works by transmitting the models with lower-bit parameters. qat works by minimizing the rounding error on the training dataset. fullyqt @xcite shows the fully quantized transformer can avoid any accuracy loss in translation quality. i-bert @xcite eliminates floating point calculation in bert inference and achieves similar accuracy to the full-precision baseline. binarybert @xcite achieves good performance by ternary weight splitting which avoids the complex and irregular loss landscape. there are some other techniques, low-rank factorization @xcite @xcite , factorizing the weight matrices which are usually low-rank into several smaller matrices by means of singular value decomposition. parameter sharing @xcite @xcite , reduces memory overhead by reusing the same parameters in multiple computations. pruning @xcite @xcite @xcite , deletes some parameters which are "useless" in network under the same model performance. zero-shot quantization. in contrast to qat, ptq relies less on the training data and does not require end-to-end training @xcite @xcite @xcite , which can be zero-shot. @xcite propose modulewise quantization error minimization and perform close to qat. smoothquant @xcite transmits the activation quantization to weights to smooth the outliers in activation and makes a large language model implemented efficiently. @xcite proposes an adaptive quantization for different modules and performs compara-bly with the full-precision generative pre-trained language models. gptq @xcite , a one-shot weight quantization method with approximate second-order information, can efficiently quantize generative pre-trained transformers with large amounts of parameters to 3 or 4 bits with negligible accuracy degradation. @xcite figures out that aggressive quantization can lead to the steep curvature of the loss landscape, while they propose a method combining layer-by-layer quantization and multivariate quadratic optimization that can improve accuracy over current ptq. nuqmm @xcite uses a non-uniform quantization method which allows for a trade-off between accuracy and compression ratio. another way to achieve zero-shot is by means of the generative adversarial technique. @xcite @xcite has shown promising performance on small-scale computer vision tasks, while it has not been applicable to the nlp realm. to the best of our knowledge, we are the first to adopt the adversarial learning approach for quantizing both the discriminative and generative plms in a zero-shot manner and provide its rigorous convergence analysis. ## methodology in this section, we will present our zero-shot sharpness-aware quantization method. network quantization. for a real-valued vector x which represents the parameter in the pre-trained full-precision model, we quantize it into x by the signed symmetric quantization: and the default experimental setting is presented in sec. 5. generative adversarial learning. in our paper we achieve zero-shot by generative adversarial learning. specifically, the generator g takes i.i.d. token samples z as input to produce synthetic data. then the generated sentence segments will be fed into the pre-trained model p and the quantized model q. intuitively we can get the output distributions, denoted as p(g(z)) and q(g(z)) respectively. during the training process of the quantized model, it aims to simulate the output distribution mse 𝒍𝒐𝒔𝒔 generator tokens pre-trained model quantized model forward pass minimize backward maximize backward figure 1: an overview of our zsaq framework. while for the generator g, it tries to maximize the divergence, so that the quantized model q can be trained to confuse even the most strict generator. we try to balance the trade-off between the generator g and the quantized model q so that we can get a well-performed quantized model without access to the original training data. the structure of our method is demonstrated in the figure 1 . and the objective function of our method is shown below: as shown in @xcite , there is a sharper loss landscape in the low-precision model than in the full-precision model. to avoid overfitting problems caused by the training process, we should consider the generalization ability of our quantized models. inspired by @xcite , we pursue a flatter landscape during the training process of the quantized model which can guarantee a better generalization performance. so we can adjust our objective function as: the sharpness function searches for the sharpest point among a neighbor centered at ω with a radius of β, which will represent the worst-case data. then our optimization process not only minimizes the divergence but also the sharpness of the divergence, i.e., getting a quantized model with both lower divergence and a flatter landscape. solving the inner maximization problem to obtain the optimal: remark. here we use ω + ϵ as the parameter in the sharpness function (4) instead of ω+ϵ, for the reason that ϵ may be so small that there will be no difference adding the perturbation. and this issue is also illustrated in @xcite . algorithm. for the objective function, we train the generator to maximize the discrepancy between the pre-trained model and the quantized model: then we train the quantized model to minimize both the discrepancy value and the sharpness of the discrepancy, i.e., the addition of d(p, q) and s(p, q), on the synthetic data produced by the generator g. while the addition comes out to be: we obtain our final algorithm by applying stochastic gradient descent on the minimization objective (7), which turns out to be sam, and applying stochastic gradient ascent on the maximization objective (6). generator g is parameterized by θ, quantized model q initially quantized from the pre-trained model p and the quantized parameter is denoted as ω, while the parameters before quantization is denoted as ω. we alternatively update the parameters ω and θ respectively. the pseudo code is shown in the algorithm 1. ## optimization in this part, we will analyse the convergence of our algorithm to solve this minimax optimization problem, which can provide a theoretical guarantee for our method. first in order to make our theoretical convergence result applicable to a broader variety of such problems, we replace the detailed loss function algorithm 1 sam-sga for zsaq input: generator g θ , quantized model q( ω), pretrained model p, learning rate η θ and η ω , and the neighborhood size β for the perturbation 1: initialize: q( ω0 ) is quantized from p by eq.(1) 2: for @xmath0,1,...,t do 3: //training quantized model by sam 4: compute the optimal perturbation ϵ * ( ω) by eq.(5); 5: calculate the gradient and update ω t+1 via eq.(7); 6: quantize ω t+1 to ωt+1 by eq .(1); 7: //training generator by sga 8: update θ t+1 by eq.( 6 ); 9: end for output: ω uniformly drawn from { ω1 , ..., ωt }. next, we restate the update rule of our algorithm 1 in our new denotation that: the first parameter can be either a floating point or an integer. it is well-defined for us to represent assumption 4.1. f (ω, θ) is differential and l-lipschitz smooth: assumption 4.2. f (ω, •) satisfies pl condition on every given ω, i.e. there exists µ > 0 such that assumption 4.3. parameter θ is restricted in a convex and bounded set with a diameter of d. assumption 4.4. the approximated gradient is unbiased and has a bounded variance, i.e., remark. these assumptions are common in the context of minimax optimization problems and can be realized in empirical practice. lemma 4.1. @xcite the quantization error is bounded by the grain, where d denotes the dimension of parameter ω: remark. the quantization error can be naturally derived from eq. ( 1 ) and the grain ∆ is decided by default quantization setting, which will not be discussed in our paper. definition 4.1. φ(ω) = max θ f (ω, θ); φ * = min ω φ(ω). and we define an ϵ-stationary point when e∥∇φ(ω)∥ @xmath1. remark. by defining the function φ, we can simplify our minimax optimization with two opposite variables into a minimization optimization problem with a single variable. when the gradient of φ(ω) diminishes to zero, we consider our algorithm converges. remark. due to the space limitation, the proofs and strongly-concave case are placed in appendix a. observing the above results, we evaluate the averaged gradients at the quantized parameter ω and we can draw the conclusion that our averaged output of the quantized models can converge to an ϵ-stationary point during o(1/ϵ 2 ) iterations under the neglect of quantization error d∆ 2 . 5 experiment 5.1 setup tasks and models. in this section, we evaluate our sam-sga method on both discriminative (i.e., bert-style) and generative (i.e., gpt-style) language models. for discriminative models, we followed many prior works @xcite ) and tested both bert base and bert large @xcite on glue benchmark @xcite ; and for generative models, we tested the optfamily @xcite (mainly opt-350m) on the glue benchmark and several language modeling tasks, i.e., wikitext2 @xcite and penn treebank (ptb) @xcite and wikitext103 @xcite . for evaluation, we report the performance with accuracy ("acc.") metric for most tasks, except the pearson correlation ("pear.") for sts-b, the matthew correlation ("mcc.") for cola, the perplexity (ppl) score for language modeling tasks. we report the averaged results over 5 random seeds to avoid stochasticity. the details of all tasks are shown in appendix a.1. implementation details. following many previous studies @xcite , we first fine-tune a full-precision model using the pretrained checkpoint from huggingface foot_1 for each task, and use the fine-tuned model as the full-precision teacher model. then, we apply the cutting-edge ptq methods, i.e., zeroquant @xcite for bert models and gptq @xcite for opt models, to quantize the fine-tuned model and use the quantized network to initialize the student model. for the generator model, we directly use the opt-350m @xcite ## conclusion in this paper, we propose a novel zero-shot quantization framework zsaq, which effectively realizes zero-shot quantization without much accuracy drop and shows promising generalization ability. we provide theoretical analysis for the minimax optimization algorithm under both nonconvex-stronglyconcave and nonconvex-pl conditions, which guarantees a convergence rate of o(1/ √ t ) for our zsaq method. our experiments show outperforming quantization results, up to +6.98 average score, and validate improved generalization ability.