Source-linked AI summary
DistilHuBERT: Speech Representation Learning by Layer-wise Distillation of Hidden-unit BERT
Heng-Jui Chang, Shu-wen Yang, Hung-yi Lee
TL;DR
Speech SSL models offer strong representations but require substantial memory and pre-training cost. The paper introduces DistilHuBERT, which layer-wise distills HuBERT hidden representations through multi-task prediction heads. It reduces HuBERT’s size by 75% and speeds inference by 73% while retaining most performance across ten tasks.
Problem
High memory and pre-training costs limit the accessibility of speech SSL models, while few studies investigate distilling speech SSL representations.
Method
DistilHuBERT uses a small student with separate prediction heads to learn the 4th, 8th, and 12th HuBERT hidden-layer representations, then removes the heads.
Results
DistilHuBERT reduces HuBERT’s size by 75% and speeds inference by 73% while retaining most performance across ten tasks.
Takeaways & Limitations
The method supports compact speech SSL representations and can be further compressed by pruning or quantization for on-device computation.
Takeaways & Limitations
The paper reports only some results because of space limitations, with a complete comparison provided on the SUPERB benchmark website.
Abstract
from arXiv · showhide
Self-supervised speech representation learning methods like wav2vec 2.0 and Hidden-unit BERT (HuBERT) leverage unlabeled speech data for pre-training and offer good representations for numerous speech processing tasks. Despite the success of these methods, they require large memory and high pre-training costs, making them inaccessible for researchers in academia and small companies. Therefore, this paper introduces DistilHuBERT, a novel multi-task learning framework to distill hidden representations from a HuBERT model directly. This method reduces HuBERT's size by 75% and 73% faster while retaining most performance in ten different tasks. Moreover, DistilHuBERT required little training time and data, opening the possibilities of pre-training personal and on-device SSL models for speech.
1. INTRODUCTION
Speech SSL methods perform well across many downstream tasks, but their memory and pre-training costs limit accessibility. DistilHuBERT addresses this gap by layer-wise distilling hidden representations from HuBERT with a multi-task framework.
- SUPERB and LeBenchmark evaluate speech SSL models across recognition, detection, semantics, and speaker identification tasks.
- High memory and pre-training costs make wav2vec 2.0 and HuBERT difficult for academia, small companies, and on-device computation.
- Few studies investigate distilling speech SSL models, motivating a framework that layer-wise distills their hidden representations.
- DistilHuBERT distills HuBERT with multi-task prediction of hidden representations and reduces HuBERT’s size by 75% and speedup by 73%.It retains most performance and requires less training time.
2. METHODS
DistilHuBERT uses a small CNN-transformer student to learn multiple HuBERT hidden-layer representations through separate prediction heads. After pre-training, the heads are removed and the frozen student supplies representations for downstream tasks.
- 2. METHODS: HuBERT has 95M to 1B parameters, consuming large memory and slowing inference.
- 2.2. DistilHuBERT: The loss combines L(4), L(8), and L(12) to train the three hidden-representation prediction tasks.
- 2.2. DistilHuBERT: The student learns multiple teacher hidden representations from shared representations using separate prediction heads in a multi-task objective.
- 2.2. DistilHuBERT: After pre-training, the prediction heads are removed, the model is frozen, and its output representations support downstream tasks.
3. EXPERIMENTS
Experiments evaluate DistilHuBERT across SUPERB tasks, model efficiency, ablations, hidden-layer selection, and distillation datasets. The results show strong performance with reduced size and faster inference, while layer choice and training data affect task outcomes.
- SUPERB Evaluation: DistilHuBERT outperformed all methods except HuBERT on the average ranking across ten SUPERB tasks.The ranking score averages each method’s ranks across the ten tasks.
- Ablation Study: Cosine similarity loss contributed more to DistilHuBERT’s improvement than teacher initialization in the ablation study.Both techniques benefited training, but removing cosine similarity loss degraded performance more.
- Layer Selection: Predicting the 4th, 8th, and 12th layers together produced less biased representations than predicting one or two selected layers.Predicting only the 4th layer was worst in IC and ASR but better in SID than predicting three layers.
- Different Distillation Datasets: Using smaller or mismatched distillation datasets caused task-dependent degradation, but DistilHuBERT remained better than many other SSL methods.The 100-hour LibriSpeech subset caused slight degradation, while AISHELL-1 performed worst in IC and ASR because of language mismatch.
4. CONCLUSION
DistilHuBERT layer-wise distills knowledge from HuBERT while retaining most of its performance and substantially reducing model size. The framework also demonstrates flexibility and generalizability, including application to more powerful SSL models.
- DistilHuBERT layer-wise distills knowledge from HuBERT and retains most of HuBERT’s performance.
- 75% smaller size leaves DistilHuBERT with only 25% of HuBERT’s size.
- DistilHuBERT’s analyses demonstrate flexibility and generalizability.
- The method can be applied to more powerful self-supervised learning models.