Effective and Efficient Mixed Precision Quantization of Speech Foundation Models

TL;DR

The paper tackles the challenge of deploying large SSL speech foundation models on resource-constrained devices by proposing a one-stage, joint mixed-precision learning and quantized parameter estimation approach. It employs NAS-based automatic mixed-precision learning with a differentiable supernet and Gumbel-Softmax DARTS to determine layerwise bit-widths, coupled with KL regularization to align quantized students with a full-precision teacher. The method achieves substantial lossless compression gains (e.g., up to 8.6x on HuBERT-large) and improves or maintains WER compared with uniform-precision and two-stage baselines, while reducing system compression time by up to $1.9\times$. Post-quantization fine-tuning (Pass 2) offers additional gains, and the framework demonstrates practical impact for on-device ASR by enabling aggressive quantization without statistically significant performance loss.

Abstract

This paper presents a novel mixed-precision quantization approach for speech foundation models that tightly integrates mixed-precision learning and quantized model parameter estimation into one single model compression stage. Experiments conducted on LibriSpeech dataset with fine-tuned wav2vec2.0-base and HuBERT-large models suggest the resulting mixed-precision quantized models increased the lossless compression ratio by factors up to 1.7x and 1.9x over the respective uniform-precision and two-stage mixed-precision quantized baselines that perform precision learning and model parameters quantization in separate and disjointed stages, while incurring no statistically word error rate (WER) increase over the 32-bit full-precision models. The system compression time of wav2vec2.0-base and HuBERT-large models is reduced by up to 1.9 and 1.5 times over the two-stage mixed-precision baselines, while both produce lower WERs. The best-performing 3.5-bit mixed-precision quantized HuBERT-large model produces a lossless compression ratio of 8.6x over the 32-bit full-precision system.

Figures (1)

• Figure 1: Diagram of joint mixed-precision learning and quantized model parameter estimation system with automatic mixed-precision learning over multiple weight-sharing systems of different quantization bit-widths. Joint mixed-precision learning and quantized model training stage (Pass 1) : all the candidate networks of different bit-widths (i.e., 2-, 4-, and 8-bit) and the mixed-precision bit-widths are simultaneously trained and updated. Post-quantization fine-tuning stage (Pass 2): parameters can be further fine-tuned while keeping the mixed-precision bit-widths frozen.