On-demand compute reduction with stochastic wav2vec 2.0

TL;DR

This paper introduces stochastic compression techniques for wav2vec 2.0 models, enabling on-demand compute reduction during inference with minimal impact on accuracy, and demonstrates effective fine-tuning for specific configurations.

Contribution

It proposes stochastic compression methods, including variable squeeze factors and pooling mechanisms, for wav2vec 2.0, allowing flexible trade-offs between accuracy and computational efficiency.

Findings

Stochastic models achieve a smooth WER-inference time trade-off.

Marginal WER degradation compared to fixed models.

Fine-tuning recovers WER with significant computational savings.

Abstract

Squeeze and Efficient Wav2vec (SEW) is a recently proposed architecture that squeezes the input to the transformer encoder for compute efficient pre-training and inference with wav2vec 2.0 (W2V2) models. In this work, we propose stochastic compression for on-demand compute reduction for W2V2 models. As opposed to using a fixed squeeze factor, we sample it uniformly during training. We further introduce query and key-value pooling mechanisms that can be applied to each transformer layer for further compression. Our results for models pre-trained on 960h Librispeech dataset and fine-tuned on 10h of transcribed data show that using the same stochastic model, we get a smooth trade-off between word error rate (WER) and inference time with only marginal WER degradation compared to the W2V2 and SEW models trained for a specific setting. We further show that we can fine-tune the same stochastically pre-trained model to a specific configuration to recover the WER difference resulting in significant computational savings on pre-training models from scratch.