Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Automatic Mixed-Precision Quantization Search of BERT

Changsheng Zhao, Ting Hua, Yilin Shen, Qian Lou, Hongxia Jin

TL;DR

Large transformer models like BERT offer strong NLP performance but are impractical on limited devices due to size and compute. The authors introduce AQ-BERT, a two-level differentiable framework that automatically assigns per-subgroup precision and prunes redundant groups via a model-size penalty, guided by a group Lasso objective. The method leverages differentiable NAS techniques (including Gumbel-Softmax) to enable end-to-end optimization without manual hyperparameter tuning, and demonstrates superior accuracy at reduced sizes across four NLP tasks, with added potential when combined with DistilBERT for extreme compression. This work enables practical deployment of compact BERT variants on mobile and edge devices with minimal performance loss, while remaining compatible with orthogonal distillation approaches.

Abstract

Pre-trained language models such as BERT have shown remarkable effectiveness in various natural language processing tasks. However, these models usually contain millions of parameters, which prevents them from practical deployment on resource-constrained devices. Knowledge distillation, Weight pruning, and Quantization are known to be the main directions in model compression. However, compact models obtained through knowledge distillation may suffer from significant accuracy drop even for a relatively small compression ratio. On the other hand, there are only a few quantization attempts that are specifically designed for natural language processing tasks. They suffer from a small compression ratio or a large error rate since manual setting on hyper-parameters is required and fine-grained subgroup-wise quantization is not supported. In this paper, we proposed an automatic mixed-precision quantization framework designed for BERT that can simultaneously conduct quantization and pruning in a subgroup-wise level. Specifically, our proposed method leverages Differentiable Neural Architecture Search to assign scale and precision for parameters in each sub-group automatically, and at the same time pruning out redundant groups of parameters. Extensive evaluations on BERT downstream tasks reveal that our proposed method outperforms baselines by providing the same performance with much smaller model size. We also show the feasibility of obtaining the extremely light-weight model by combining our solution with orthogonal methods such as DistilBERT.

Automatic Mixed-Precision Quantization Search of BERT

TL;DR

Large transformer models like BERT offer strong NLP performance but are impractical on limited devices due to size and compute. The authors introduce AQ-BERT, a two-level differentiable framework that automatically assigns per-subgroup precision and prunes redundant groups via a model-size penalty, guided by a group Lasso objective. The method leverages differentiable NAS techniques (including Gumbel-Softmax) to enable end-to-end optimization without manual hyperparameter tuning, and demonstrates superior accuracy at reduced sizes across four NLP tasks, with added potential when combined with DistilBERT for extreme compression. This work enables practical deployment of compact BERT variants on mobile and edge devices with minimal performance loss, while remaining compatible with orthogonal distillation approaches.

Abstract

Pre-trained language models such as BERT have shown remarkable effectiveness in various natural language processing tasks. However, these models usually contain millions of parameters, which prevents them from practical deployment on resource-constrained devices. Knowledge distillation, Weight pruning, and Quantization are known to be the main directions in model compression. However, compact models obtained through knowledge distillation may suffer from significant accuracy drop even for a relatively small compression ratio. On the other hand, there are only a few quantization attempts that are specifically designed for natural language processing tasks. They suffer from a small compression ratio or a large error rate since manual setting on hyper-parameters is required and fine-grained subgroup-wise quantization is not supported. In this paper, we proposed an automatic mixed-precision quantization framework designed for BERT that can simultaneously conduct quantization and pruning in a subgroup-wise level. Specifically, our proposed method leverages Differentiable Neural Architecture Search to assign scale and precision for parameters in each sub-group automatically, and at the same time pruning out redundant groups of parameters. Extensive evaluations on BERT downstream tasks reveal that our proposed method outperforms baselines by providing the same performance with much smaller model size. We also show the feasibility of obtaining the extremely light-weight model by combining our solution with orthogonal methods such as DistilBERT.
Paper Structure (27 sections, 12 equations, 1 figure, 3 tables, 1 algorithm)

This paper contains 27 sections, 12 equations, 1 figure, 3 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Proposal of an united framework to achieve automatic parameter search.
  3. Proposal of a novel objective function that can compress model to the desirable size, and enables pruning and quantization simultaneously.
  4. Provide efficient solutions to optimize the parameters for the proposed framework.
  5. Extensive experimental validation on various NLP tasks.
  6. Related Work
  7. General Model Compression
  8. Transformers Compression
  9. Network Architecture Search
  10. Methods
  11. Framework
  12. Objective Function
  13. Quantization Process
  14. Continuous Relaxation
  15. Optimization Process
  16. ...and 12 more sections

Figures (1)

  • Figure 1: Framework. The central part shows the idea of two-stage optimization. The left part is the illustration of the inner training network, while the right part is an example of the super network that controls bit-assignment. In the left part, each node represents a matrix (a group of neurons), which we call a "sub-group" in our paper. Each sub-group has its own quantization range in the mixed-precision setting. As the example shown in the right part, a sub-group has three choices of bit-assignment: 0-bit, 2-bit, and 4-bit. And each such assignment is associated with a probability of being selected.