Vision transformer model quantization method based on attention head importance guidance
By optimizing the quantization process of the Vision Transformer model through a quantization method guided by the importance of attention heads, the problem of accuracy loss in low-bit cases is solved, and efficient deployment and accuracy preservation are achieved on edge devices.
Patent Information
- Application Number
- CN202610226318.4
- Authority / Receiving Office
- CN · China
- Patent Type
- Applications(China)
- Current Assignee / Owner
- Filing Date
- 2026-02-26
- Publication Date
- 2026-06-23
Smart Images
Figure CN122263982A_ABST
Abstract
Description
Technical Field
[0001] This invention belongs to the field of low-bit quantization of deep learning models, specifically involving a quantization method for the Vision Transformer model based on attention head importance guidance. Background Technology
[0002] With the development of deep neural networks, the demand for lightweight models in edge devices is increasing. Model compression technology has become a research focus for model deployment on edge devices. Vision Transformer performs well in computer vision tasks, but its large number of model parameters and high computational overhead make it difficult to deploy widely on resource-constrained edge devices.
[0003] Low-bit model quantization, as one of the mainstream model compression techniques, reduces the number of model parameters and computational overhead by mapping data from high precision to low precision. The main quantization methods are divided into two types: (1) Post-Training Quantization (PTQ) is a pre-trained full-precision model. It collects the distribution of its weights or activation values and determines the scaling factor and zero point for each layer based on statistical information. It can be directly converted into a low-precision model without retraining.
[0004] (2) Quantization-Aware Training (QAT) simulates the quantization process during the training phase. A pre-trained full-precision model is prepared, and pseudo-quantization operators are inserted at the output positions of weights and activation values that require quantization operations (such as convolution and fully connected layers) to simulate the quantization effect during forward propagation, but a pass-through estimator is used to pass the gradient during backpropagation. During the learning process, the model actively adapts to the error brought about by quantization and improves the accuracy of the quantized model.
[0005] While QAT offers high accuracy in quantization methods for Vision Transformers, quantization occurs during training, requiring model retraining, which is time-consuming. This method is also limited when the original dataset is inaccessible. Traditional PTQ, requiring no additional training and being simpler, has become the mainstream quantization method in industry and a research hotspot in academia in recent years. However, PTQ, after determining quantization parameters layer by layer independently, does not consider the propagation and accumulation of quantization errors between layers, leading to significant accuracy loss (especially in low-bit cases). These issues have become the core bottleneck for its further application.
[0006] To address the severe precision loss in PTQ at low bit depths, recent research has employed weight perturbation and output reconstruction methods (such as AdaRound and QDrop) to replace minimizing parameter error with minimizing task loss error. Taylor expansion analysis is used to analyze loss changes and learn a better quantization rounding strategy to replace traditional rounding. In other words, during quantization, the goal is not only to minimize weight error but also to make the output of the entire quantized network layer as close as possible to the original full-precision network output.
[0007] In the Transformer model structure, many attention heads are relatively redundant and have a limited impact on the results. However, in the reconstruction-based quantization method for Vision Transformer, all attention heads are treated equally when optimizing for quantization error, resulting in the model not being fully optimized. Summary of the Invention
[0008] Purpose of the invention: To solve the above problems, this invention aims to propose a Vision Transformer model quantization method based on the importance of attention heads.
[0009] Technical Solution: A quantization method for Vision Transformer models based on attention head importance guidance, applied to low-bit quantization of deep learning models. This method performs the following operations on a fully-precision pre-trained Vision Transformer model: S1. Calculate the importance of each attention head in the Attention module of the model based on Fisher information. Use the empirical Fisher information with diagonal approximation for calculation, and then obtain the importance score of each attention head based on Fisher information. S2. Initialize the quantization parameters. Use a channel-level uniform quantizer to implement weight quantization and a hierarchical uniform quantizer to implement activation quantization. At the same time, initialize the quantization weights through the AdaRound strategy. S3. Calculate the weighted loss based on the importance scores of each attention head, assign weights matching the importance scores to different attention heads, and then calculate the mean squared error loss by combining the full-precision and analog quantization outputs to obtain the overall loss after weighting the attention heads. S4. Iteratively optimize the quantization weights, calculate the gradient of the loss with respect to the quantization parameters through backpropagation, update the parameters to minimize the weighted loss function, and obtain the quantized Vision Transformer model after the iteration is completed.
[0010] Further, step S2 includes: S21. Register a hook in the VIT Attention module and save the output of each attention head; S22. Calculate the empirical Fisher information for the diagonal approximation of each attention head. The calculation formula is as follows:
[0011] In the formula, Indicates the given input and parameters Time prediction output The probability, Indicates the size of the calibration dataset. These are model probability parameters; S23: Calculate the attention importance score based on the obtained Fisher information. The calculation formula is as follows:
[0012] In the formula, Indicates the first The importance score of each attention head For the first The set of parameters corresponding to each attention head.
[0013] Furthermore, the output of the h-th attention head saved by hook in step S21 The calculation formula is:
[0014] In the formula, , , They are the first The query, key, and value matrix of the size. It is the dimension of the key. The shape is [head_dim], where It is the batch size. It is the number of tokens.
[0015] Furthermore, the process of initializing the quantization parameters in step S3 specifically includes: S31. Input the calibration dataset into the full-precision model, obtain the original full-precision output of the target block as the benchmark output for subsequent reconstruction, and use it to calculate the quantization error; S32. Use channel-level uniform quantizer and hierarchical uniform quantizer to initialize the scaling factor, zero point and other parameters required for quantization, and set the hyperparameters for quantization reconstruction. S33. Using the AdaRound strategy, the weight quantization is treated as a perturbation of the full-precision weights. The quantization weights are initialized using the following formula: ,in These are the quantized model weights. For full-precision model weights, This is a weighted perturbation.
[0016] Furthermore, the calculation steps for the attention head score weighted loss in step S4 are as follows: S41. Assign different weights to each attention head based on its score, as calculated below:
[0017] In the formula, It is the first The weight of each attention head, For the total number of attention heads, For the first The importance score of each attention head; S42. Calculate the loss of a single attention head. Use the mean square error of the full-precision output and the analog quantization error as the loss function. The calculation formula is as follows:
[0018] In the formula, For the first The loss of attention, To calibrate the sample size of the dataset, For the first The output of the attention head simulation quantization For the first Full-precision output from a single attention head; S43. Calculate the overall weighted loss using the weights of each attention point. The calculation formula is as follows:
[0019] In the formula, The overall loss after weighting by attention.
[0020] Furthermore, the specific sub-steps for iteratively optimizing the quantization weights in step S5 are as follows: S51. Calculate the gradient of the weighted loss with respect to the quantization parameters through backpropagation. The quantization parameters include the weight rounding value optimized by the AdaRound policy and the quantization scaling factor. S52. Based on gradient-based quantization parameter updates, minimize the weighted loss function to prioritize the optimization of attention heads with higher importance; S53. Repeat steps S51-S52 until the iteration ends to obtain the fully quantized Vision Transformer model. Perform inference to verify the accuracy of the model.
[0021] Furthermore, the method does not require retraining the model using the original dataset; it only optimizes the quantization parameters and reconstructs the model by calibrating the dataset.
[0022] This method is applicable to low-bit quantization of various Transformer-based models, can be deployed on resource-constrained edge devices, and can be used to achieve computer vision tasks such as image classification and object detection.
[0023] Furthermore, the fully-precision pre-trained Vision Transformer model described in this invention is designed for edge device image classification and object detection tasks, and also includes one or more of ViT-Base, ViT-Small, Deit-Base, and Swin-Base.
[0024] Beneficial effects: Compared with the prior art, the substantial features and significant advancements of this invention include: (1) This invention is a reconstruction quantization method based on the importance of attention heads. Combining the characteristics of the multi-head attention mechanism in VisionTransformer, it can prioritize the optimization of more important parameters during quantization reconstruction. Reconstruction quantization is more targeted when optimizing VisionTransformer-type models.
[0025] (2) The present invention adopts the reconstruction quantization method, which significantly reduces the computational overhead and memory usage of the Vision Transformer model by mapping the model from high bits to low bits, and maintains high accuracy without requiring retraining on the original dataset.
[0026] (3) The quantitative reconstruction optimization of the multi-head attention mechanism in this invention can be widely applied to various Transformer-based models. Attached Figure Description
[0027] To more clearly illustrate the technical methods and advantages of the present invention, the following figures provide a detailed description of the invention's content and process.
[0028] Figure 1 This is a flowchart illustrating the workflow of the Vision Transformer quantization method based on attention head importance proposed in this invention. Figure 2 This is a schematic diagram of the attention importance weighted loss module in this invention; Figure 3 This is a schematic diagram of the weighted loss structure of the Vision Transformer quantization method based on attention head importance proposed in this invention. Detailed Implementation
[0029] The technical content of the invention will be further described below with reference to specific embodiments. These embodiments are for explaining and illustrating the invention, and not for limiting the scope of the invention. All other embodiments obtained by those skilled in the art based on the embodiments of the invention without creative effort are within the protection scope of the invention.
[0030] This embodiment provides a Vision Transformer model quantization method based on attention head importance guidance, specifically for image classification and object detection tasks on edge devices. In a specific deployment scenario, a pre-trained full-precision model can be reconstructed and quantized to obtain a low-bit model, which can then be deployed to edge devices. This maintains model accuracy with low computational overhead and memory usage.
[0031] Specifically, the implementation of the method described in this invention includes the following steps: S1: Train a full-precision model based on the dataset; 1) Dataset: The dataset used in this embodiment is the ImageNet dataset. ImageNet is a large-scale image dataset containing approximately 14 million images, of which the training set (train) contains approximately 1.2 million images, the validation set (val) contains 50,000 images, and the test set (test) contains 50,000 images. The ImageNet dataset contains 1,000 categories, covering many common items such as animals, plants, furniture, tools, and electronic devices.
[0032] 2) Model selection: The experimental models selected included several types of VIT models: VIT-Base, VIT-Small, Deit-Base, and Swin-Base.
[0033] S2: Calculate the Fisher information for each attention head in the model, and further calculate the importance score; The method for this step is as follows: Figure 2 As shown, the specific process is as follows: (1) In the Attention module of VIT A hook is registered to store the output of each attention head. The loss is computed on the calibration set and backpropagated, with the hook capturing the output gradient of each head. The formula is:
[0034] in: They are the first The query, key, and value matrix of the size. It is the dimension of the key (usually equal to head_dim). The shape is [head_dim], where It is the batch size. It is the number of tokens.
[0035] (2) Calculate the Fisher information for each attention head. Since the computational cost of the complete Fisher information matrix is too high, the empirical Fisher information, which is approximated by a diagonal, is used. The Fisher information matrix measures the amount of information about the observed data X carried by the probability model parameter θ, and is calculated as follows:
[0036] (3) Based on the Fisher information obtained, calculate the attention importance score, as shown in the following formula:
[0037] S3. Initialize the quantization parameters. The specific steps are as follows: (1) Input the calibration dataset into the full-precision model and obtain the output of each sample of the target block as the “benchmark output” for subsequent reconstruction, which is used to calculate the quantization error.
[0038] (2) Use channel-level uniform quantizer (weighted quantization) and hierarchical uniform quantizer (activation quantization) to initialize the scaling factor, zero point and other parameters required for quantization, and set the hyperparameters for quantization reconstruction.
[0039] (3) The AdaRound strategy is used to initialize the quantization weights. This method is different from the traditional "rounding". The steps for calculating the attention head score weighted loss in S4 by optimizing the weight rounding method are as follows: S4: Calculation of attention head score weighted loss; (1) Assign different weights to each attention head based on its score, as follows:
[0040] (2) The mean square error of the full-precision output and the analog quantization error output (quantized first and then dequantized) is used as the loss function:
[0041] (3) Use the obtained weights to weight the loss of each attention point:
[0042] S5. Iteratively optimize the quantization parameters.
[0043] (1) Calculate the gradient of the loss with respect to the quantization parameters (AdaRound optimized weight rounding value, quantization scaling factor) through backpropagation.
[0044] (2) Update the parameters and minimize the weighted loss function to optimize the attention heads with higher importance.
[0045] (3) After the iteration is completed, the quantized model is obtained, and the accuracy is verified by reasoning.
[0046] To further verify the advantages of this invention, the head-weighted loss quantization method (activation quantization uses the MLP reconstruction method in APHQ-VIT) is compared with other quantization methods. The results are shown in Table 1. This invention (HeadAware Quantization, HAQ) improves the performance of all models, achieving higher classification accuracy on the ImageNet dataset after application to each VIT class model. Here, W / A represents the quantization bit width of the weights and activations.
[0047] Table 1 Evaluation results of the model quantification method based on attention importance guidance
[0048] In summary, the Vision Transformer model quantization method based on attention head importance provided in this embodiment of the invention fully utilizes the characteristics of the multi-head attention mechanism in Vision Transformer, prioritizing the quantization of more important attention heads, thereby further improving the model's performance with low bit complexity and leveraging the AdaRound technique.
[0049] The embodiments of the present invention have been described in detail above with reference to the accompanying drawings, but the present invention is not limited to the described embodiments. For those skilled in the art, various changes, modifications, substitutions, and variations can be made to these embodiments, including components, without departing from the principles and spirit of the present invention, and these variations still fall within the protection scope of the present invention.
Claims
1. A Vision Transformer model quantization method based on attention head importance guidance, applied to deep learning model low-bit quantization, characterized in that, This method performs the following operations on a fully pre-trained Vision Transformer model: S1. Calculate the importance of each attention head in the Attention module of the model based on Fisher information. Use the empirical Fisher information with diagonal approximation for calculation, and then obtain the importance score of each attention head based on Fisher information. S2. Initialize the quantization parameters. Use a channel-level uniform quantizer to implement weight quantization and a hierarchical uniform quantizer to implement activation quantization. At the same time, initialize the quantization weights through the AdaRound strategy. S3. Calculate the weighted loss based on the importance scores of each attention head, assign weights matching the importance scores to different attention heads, and then calculate the mean squared error loss by combining the full-precision and analog quantization outputs to obtain the overall loss after weighting the attention heads. S4. Iteratively optimize the quantization weights, calculate the gradient of the loss with respect to the quantization parameters through backpropagation, update the parameters to minimize the weighted loss function, and obtain the quantized Vision Transformer model after the iteration is completed.
2. The Vision Transformer model quantization method based on attention head importance guidance according to claim 1, characterized in that, Step S2 includes: S21. Register a hook in the VIT Attention module and save the output of each attention head; S22. Calculate the empirical Fisher information for the diagonal approximation of each attention head. The calculation formula is as follows: In the formula, Indicates the given input and parameters Time prediction output The probability, Indicates the size of the calibration dataset. These are model probability parameters; S23: Calculate the attention importance score based on the obtained Fisher information. The calculation formula is as follows: In the formula, Indicates the first The importance score of each attention head For the first The set of parameters corresponding to each attention head.
3. The Vision Transformer model quantization method based on attention head importance guidance according to claim 2, characterized in that, In step S21, the output of the h-th attention head is saved via hook. The calculation formula is: In the formula, , , They are the first The query, key, and value matrix of the size. It is the dimension of the key. The shape is [head_dim], where It is the batch size. It is the number of tokens.
4. The Vision Transformer model quantization method based on attention head importance guidance according to claim 1, characterized in that, The process of initializing the quantization parameters in step S3 specifically includes: S31. Input the calibration dataset into the full-precision model, obtain the original full-precision output of the target block as the benchmark output for subsequent reconstruction, and use it to calculate the quantization error; S32. Use channel-level uniform quantizer and hierarchical uniform quantizer to initialize the scaling factor, zero point and other parameters required for quantization, and set the hyperparameters for quantization reconstruction. S33. Using the AdaRound strategy, the weight quantization is treated as a perturbation of the full-precision weights. The quantization weights are initialized using the following formula: ,in These are the quantized model weights. For full-precision model weights, This is a weighted perturbation.
5. The Vision Transformer model quantization method based on attention head importance guidance according to claim 1, characterized in that, The steps for calculating the attention head score weighted loss in step S4 are as follows: S41. Assign different weights to each attention head based on its score, as calculated below: In the formula, It is the first The weight of each attention head, For the total number of attention heads, For the first The importance score of each attention head; S42. Calculate the loss of a single attention head. Use the mean square error of the full-precision output and the analog quantization error as the loss function. The calculation formula is as follows: In the formula, For the first The loss of attention, To calibrate the sample size of the dataset, For the first The output of the attention head simulation quantization For the first Full-precision output from a single attention head; S43. Calculate the overall weighted loss using the weights of each attention point. The calculation formula is as follows: In the formula, The overall loss after weighting by attention.
6. The Vision Transformer model quantization method based on attention head importance guidance according to claim 1, characterized in that, The specific sub-steps for iteratively optimizing the quantization weights in step S5 are as follows: S51. Calculate the gradient of the weighted loss with respect to the quantization parameters through backpropagation. The quantization parameters include the weight rounding value optimized by the AdaRound policy and the quantization scaling factor. S52. Based on gradient-based quantization parameter updates, minimize the weighted loss function to prioritize the optimization of attention heads with higher importance; S53. Repeat steps S51-S52 until the iteration ends to obtain the fully quantized Vision Transformer model. Perform inference to verify the accuracy of the model.
7. The Vision Transformer model quantization method based on attention head importance guidance according to any one of claims 1-6, characterized in that, The method does not require retraining the model using the original dataset; it optimizes the quantization parameters and reconstructs the model solely by calibrating the dataset.
8. The Vision Transformer model quantization method based on attention head importance guidance according to any one of claims 1-6, characterized in that, This method is applicable to low-bit quantization of various Transformer-based models, can be deployed on resource-constrained edge devices, and can be used to achieve computer vision tasks such as image classification and object detection.
9. The Vision Transformer model quantization method based on attention head importance guidance according to claim 1, characterized in that, The fully-precision pre-trained Vision Transformer model is designed for image classification and object detection tasks on edge devices, and may include one or more of ViT-Base, ViT-Small, Deit-Base, and Swin-Base.