Shrinking Bigfoot: Reducing wav2vec 2.0 footprint

1. Introduction

• The largest wav2vec 2.0 version has 317 million parameters, making real-time inference inefficient in production.The framework combines self-supervised masking pre-training with tuning and achieves strong LibriSpeech WERs using only 10 minutes of labeled data.Source: §1 · p. 1
• Wav2vec 2.0’s self-supervised learning supports training an end-to-end ASR system with little labeled data, including for a low-resource language or specific domain.The passage relates wav2vec 2.0 to BERT’s masked pre-training and transformer architecture.Source: §1 · p. 1
• The authors explore compression methods borrowed from large language models to improve wav2vec 2.0’s applicability to production.They describe the work as the first to compress wav2vec 2.0, to their knowledge.Source: §1 · p. 1
• The distilled student is at least 2 times faster on both CPU and GPU and 4.8 times smaller than the original, with a 7% WER loss; quantization yields a model 3.6 times smaller with a 0.1% WER loss.Source: §1 · p. 1

2. Related work

• Related work on ASR compression motivates adapting language-model compression methods to wav2vec.The authors frame their review around ASR compression and language-model methods they wish to explore.Source: §2 · p. 1

2.1. Model compression in ASR systems

• ASR compression methods vary by architecture; prior work includes recurrent-model factorization and distillation, with reported accuracy reductions in some studies and improvements in others.Work on transformer ASR also explored parameter sharing and low-rank transformers, reporting accuracy and size or speed gains.Source: §2.1 · pp. 1–2

2.2. Classic compression methods in language modeling

• Language-model compression includes knowledge distillation, quantization, and pruning; distillation trains a smaller student using a teacher model.Prior ASR work also uses distillation to improve accuracy or reduce model size, and DistilBERT demonstrates size and speed gains while preserving language understanding.Source: §2.2 · p. 2
• Quantization reduces weight precision and can quantize activations statically or dynamically, while pruning removes small weights likely to be redundant.The authors use PyTorch rather than Distiller for their quantization study because Distiller simulates quantization while retaining 32-bit floating-point weights.Source: §2.2 · p. 2
• The preliminary pruning experiment achieved 23% sparsity while maintaining the original model’s WER, but did not significantly improve inference speed.The authors pruned weights according to layer-specific sensitivity thresholds.Source: §2.2 · p. 2

3. Wav2vec

• Wav2vec 2.0 masks parts of audio representations and trains a transformer to fill the gaps, learning mainly from unlabeled data before fine-tuning on a small labeled dataset.The model has CNN and transformer components, both considered for compression.Source: §3 · p. 2
• The model’s two components considered for compression are its CNN layers and transformer layers.Source: §3 · p. 2
• The study evaluates how knowledge distillation and quantization affect the size and inference speed of a fine-tuned wav2vec 2.0 model.Source: §3 · p. 2

4. Knowledge distillation for wav2vec 2.0

• The distillation approach adapts the BERT-to-DistilBERT strategy, initializing the student from alternating teacher layers and increasing then decreasing the learning rate.The authors reduce transformer layers because the seven CNN layers are not a substantial computational bottleneck relative to the transformer.Source: §4 · p. 2

4.1. Knowledge distillation loss

• The distillation loss is the KL divergence between teacher and student distributions over projected speech tokens, which the student optimizes to match the teacher.The distributions have dimensions N × M, with N the representation length and M the number of tokens.Source: §4.1 · p. 2

4.2. The architecture of the student model

• The student differs from original wav2vec 2.0 only in transformer-layer count: the largest version has seven convolution and 24 transformer layers, while student variants use 4, 6, 8, 10, or 12 transformer layers.Source: §4.2 · pp. 2–3

4.3. Objective function

• The training objective adds the distillation loss to an L2 feature penalty on the final convolution-layer output; cosine embedding loss did not improve student performance.The feature penalty serves as a regularizer for convolution layers.Source: §4.3 · p. 3

4.4. Training techniques

• Among the training changes tried, DistilBERT’s initialization strategy and training on more data significantly improved student performance.The authors also tried different optimizers, schedules, and objective functions.Source: §4.4 · p. 3

5. Quantization for wav2vec 2.0

• PyTorch’s available quantization techniques are applied to wav2vec 2.0.The section discusses both static and dynamic quantization.Source: §5 · p. 3

5.1. Static quantization

• Static quantization requires calibration for activations, but PyTorch did not support quantizing GELU as of March 2021.Dequantizing before GELU and requantizing afterward is a possible workaround that causes some performance degradation.Source: §5.1 · p. 3

5.2. Dynamic quantization

• Dynamic quantization quantizes weights and activations during each forward pass, but tensors entering multi-head attention must remain unquantized.The authors save MHA weights and biases before inference to avoid a costly linear unpack operation for every sample.Source: §5.2 · p. 3

6. Experiments

• The compressed model started from wav2vec 2.0 trained on 960 hours of LibriSpeech data and fine-tuned on the same labeled hours.The authors did not attempt compression before fine-tuning; student outputs were decoded with Viterbi.Source: §6 · p. 3
• Knowledge distillation used Adam with weight decay and a learning rate that warmed up for two epochs before decreasing linearly.Training used eps = 10^-6, betas = (0.9, 0.98), weight decay 0.01, and an initial learning rate of 2.5 × 10^-5.Source: §6 · p. 3
• Dynamic quantization was applied to wav2vec 2.0 linear layers, and compression methods were evaluated on LibriSpeech train-clean-100.Training used four Nvidia V100 GPUs, batch size one per GPU, 16-bit arithmetic, and seed 42.Source: §6 · p. 3

7. Results and discussion

• 7% WER loss accompanied a distilled model at least 2 times faster on CPU and GPU and 4.8 times smaller; quantization was 3.6 times smaller but not significantly faster.The quantized model maintained the original WER with a slight loss; only CPU inference time was reported because PyTorch did not support GPU quantization.Source: §7 · p. 3, Table 1 · p. 4

7.1. Knowledge distillation using different settings

• Fewer student transformer layers reduced inference time but increased WER in experiments trained on train-clean-100 and evaluated on a single GPU.Each student model was trained for ten epochs.Source: §7.1 · p. 3, Table 2 · p. 4
• Reducing convolution layers from seven to four made GPU inference 0.7 seconds faster but CPU inference slower.The authors attribute the CPU slowdown to replacing convolution layers with slower max-pooling operations.Source: Table 3 · p. 4, §7.1 · p. 4
• DistilBERT’s alternating-layer initialization significantly outperformed initialization from the last few layers from the beginning of training.Figure 1 compares students initialized with different strategies and then trained using knowledge distillation.Source: Figure 1 · p. 4, §7.1 · p. 4
• The student trained with 360 hours of data achieved better WER than the student trained with 100 hours.Figure 2 compares students trained with different amounts of data.Source: Figure 2 · p. 4, §7.1 · p. 4

• Wav2vec models are presented as valuable for democratizing ASR, and the work explores reducing memory use and speeding models for deployment.Source: §8 · p. 4
• 72% memory-footprint reduction from quantization is significant for reducing the costs of machines running the models.Source: §8 · p. 4
• Knowledge distillation is the most effective method explored for saving energy by lowering CPU or GPU inference time.Source: §8 · p. 4
• The distilled student model is at least two times faster than the original on both CPU and GPU.Source: §8 · p. 4
• The authors see room to further improve compressed-model performance by using more data in distillation training.Source: §8 · p. 4