Intern talk: Distilling Self-Supervised-Learning-Based Speech Quality Assessment into Compact Models
This research explores the distillation and pruning of large, self-supervised speech quality assessment models into compact and efficient versions. Starting with the high-performing but large XLSR-SQA model, the work details a process of knowledge distillation using a teacher-student framework with a diverse, on-the-fly generated dataset. The resulting compact models successfully close over half the performance gap to the teacher, making them suitable for on-device and production applications where model size is a critical constraint.
Motivation for Computational Speech Quality Assessment
Computational models for Speech Quality Assessment (SQA) are essential for scenarios where human-based listening experiments are too costly or time-consuming. The standard approach involves collecting Mean Opinion Scores (MOS) from human raters on a 1-5 scale. A non-intrusive computational model aims to predict this MOS using only the degraded speech signal as input, without a clean reference.
Applications for such models include:
- Model selection and hyperparameter tuning during development.
- Assessing the quality of datasets.
- Serving as a cost function in unsupervised learning.
- Personalization of on-device speech enhancement.
- In-production quality monitoring.
Baseline Model Evaluation
An extensive evaluation was conducted across numerous datasets, including public telephone/VoIP data (PSTN, Niska), Text-to-Speech (TTS) data, and internal data from challenges like the DNS, PLC, and Signal Improvement challenges.
Several state-of-the-art models were benchmarked:
- DNSMOS: A CNN trained on DNS challenge data.
- Niska: A CNN-Transformer model trained on VoIP data.
- TorchAudio-Squim: An RNN-Transformer model trained to predict full-reference metrics (PESK, SI-SDR).
- UT-MOS: A wav2vec 2-based model trained on TTS data.
- XLSR-SQA: A model using embeddings from a large, pre-trained wav2vec XLSR model, which won the Conferencing Speech Challenge.
The evaluation revealed that XLSR-SQA demonstrated the best generalization, performing well even on tasks and challenge data it was not explicitly trained on. Conversely, predicting the quality of TTS and voice conversion proved challenging for most models, with the distilled models sometimes amplifying the teacher’s errors.
Improving and Compacting the XLSR Model
The baseline XLSR model, while powerful, is very large (1.8 GB), making it impractical for many applications. The research focused on two main goals: slightly improving its performance and drastically reducing its size.
-
Performance Improvement: By incorporating a wider range of training data and employing techniques like random cropping and dataset bias-aware training, the model’s Pearson correlation coefficient (PCC) with subjective scores was improved from a weighted mean of 0.78 to 0.81.
-
Model Size Reduction: The primary focus was on making the model compact without sacrificing too much performance. Two techniques were explored:
- Knowledge Distillation: A teacher-student approach was implemented. The large, improved XLSR model acted as the “teacher,” providing pseudo-labels for a vast, unlabeled dataset. This dataset was generated on-the-fly with a wide array of degradations (noise, various TTS models, codecs, enhancement artifacts) to ensure diversity. A compact “student” model, with a CNN-Transformer architecture, was trained on these pseudo-labels, sometimes mixed with a small percentage of original labeled data.
- Pruning: As an alternative, pruning techniques were applied to the full-sized model. Importance-based pruning (which approximates a weight’s importance based on its impact on the loss function) was effective, particularly at larger model sizes. Standard magnitude-based pruning was not successful.
Key Results
- Distillation: The distillation process was highly effective. A distilled model with only 4 million parameters (a >100x reduction from the teacher’s ~450M parameters) closed more than half of the performance gap between the teacher model and a baseline student trained only on labeled data. As model size increases, distilled models continue to improve, whereas models trained only on labeled data tend to overfit.
- Pruning: Importance-based pruning also successfully reduced model size while retaining performance, bridging the gap to the unpruned model at larger sizes.
- Final Outcome: The research successfully developed schemes to create compact, efficient, and high-performing speech quality assessment models, providing a menu of options to balance the trade-off between computational cost and prediction accuracy. The final models were exported to ONNX format to facilitate their use in production environments.