Automatic Mixed-Precision Quantization Search of BERT

TL;DR

This paper introduces an automatic mixed-precision quantization and pruning framework for BERT that optimizes model size and performance, leveraging neural architecture search for fine-grained parameter management.

Contribution

It presents a novel differentiable neural architecture search-based method for automatic mixed-precision quantization and subgroup-wise pruning of BERT models.

Findings

Outperforms baseline models in downstream tasks.

Achieves significant model size reduction with maintained accuracy.

Enables extremely lightweight models when combined with other methods.

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.