1. 背景

近年来，随着GPU算力的不断提升，大型自监督预训练模型在自然语言处理这项任务中大放异彩，从GPT，到BERT，再到GPT-3，SOTA被不断的刷新。这类模型在预训练过程中，仅使用无标注数据进行学习，便能够对数据本身进行分析，学习通用且鲁棒的高阶特征。预训练之后的模型可以被视为一个极为优秀的特征抽取器，适合应用到诸多下游任务。然而，预训练模型在语音这一领域却迟迟没有突破。直到Wav2vec 2.0的出现，才第一次有语音预训练模型达到了SOTA表现。在此之后，越来越多的语音预训练模型才开始涌现，例如HuBERT，WavLM等。这些模型不仅在ASR任务上有着卓越的表现，在其余各项语音任务（语音分离，说话人验证等）也取得了SOTA结果，大有一统江湖的趋势。
在使用预训练模型应用到下游任务时，我们仅仅需要在原先模型的基础上稍加改动，便可以用任务相关的损失函数进行微调（finetune）。例如，在预训练好的HuBERT模型的最后加上一层线性层（linear layer），即可使用CTC损失函数对进行ASR任务的微调。Fairseq官方就提供了在LibriSpeech上微调后的HuBERT+CTC模型，使用了的建模单元为英文字母。尽管该模型的识别效果已经非常优异，但这并不能满足大家的需求：有的人需要在研究中使用不同的建模单元（比如BPE等），或是需要更为强大的模型结构（比如Transducer结构）。为了解决大家的这一需求，我们今天将介绍如何在icefall中对预训练模型使用pruned RNNT loss进行微调。目前，icefall仅支持对HuBERT模型的微调，但如果之后大家对别的预训练模型（Wav2vec 2.0， WavLM等）有支持的需求，我们也会尽快在icefall中实现，欢迎大家向我们留言~

2. HuBERT 简介

2.1 语音自监督的困境

在这一章，我们对HuBERT模型进行一个简单的介绍。开始前，咱们先来聊聊他的前辈Wav2vec2.0。在Wav2vec2.0之前，语音的预训练一直效果不佳，归根结底是语音的特征表达相较于图片和文本要复杂的多：
  1. 语音中的标签不是唯一的。在图片分类中，一张图片的标签是固定的，然而语音则没有确定的标签
  2. 语音是一个连续的特征空间，不像文本一样存在确定的词典。因此，在预训练中常见的损失函数（例如分类）对语音并不能直接使用。

2.2 Wav2vec 2.0的问世

为了解决语音特征连续化的这一问题，Wav2vec 2.0巧妙地采用乘积量化（product quantization），将语音的特征向量进行量化。这一操作将本来连续的向量空间转化为了有限的离散空间，大大提升了训练得到的特征的鲁棒性。再此之上，使用对比学习（Contrastive learning）的方式，进行遮罩语言模型（Masked language modelling)对模型进行训练。为了使量化的过程可导，Wav2vec 2.0使用了Gumbel softmax选择乘积量化中的不同codebook，并且引入了一个辅助损失函数，尽可能增加量化器在选择codebook时的多样性。

2.3 HuBERT的改进

这一通操作后，Wav2vec 2.0在ASR任务中取得了SOTA表现，成为了第一个达成这项成就的语音预训练模型。但是由于使用对比学习进行训练，需要精确的设计正负样本的采样，再加上量化器的存在，训练时的损失函数较为复杂，使得Wav2vec2.0并不稳定，并且在一些非ASR的下游任务中表现不佳。那么，如何在不量化的同时，将语音特征的空间离散化呢？换一种思路，如果我们直接将整个特征空间按照区域进行区分，是不是就可以将预训练任务转化为分类任务？我相信这时，大家很容易就会想到聚类（clustering）这一办法。没错，我们今天的主角HuBERT正是通过聚类（K-Means）的办法，对本没有标签的高维的语音特征向量贴上伪标签（pseudo label），再套用BERT的算法进行预训练。
HuBERT模型可以分为如下两个部分：CNN Encoder和Transformer。CNN Encoder的输入为原始的音频，是一个一维的向量w。通过CNN Encoder的降采样和升维，这个一维向量被转化为输出频率接近50Hz的二维音频特征X=[x1,...,xT]，每个特征之间相隔的时间为1/50秒，即20ms。X所对应的标签序列记为Z=[z1,...,zT]。在 X 进入Transformer前，会先进行随机的遮罩。具体的遮罩办法借鉴了Wav2vec 2.0的思路：先随机选取一些起始点，再从这些起始点开始，用一个相同的Embedding遮罩之后连续M个时间点的 xt，得到遮罩后的序列 X'。X' 会作为Transformer的输入，用其强大的序列建模能力，对每一个输入特征x't 都做出一次预测，希望无论是被遮罩还是未被遮罩的输入 x't ，都能够预测出正确的 zt 。不难看出，这种遮罩训练的方式非常类似于BERT(其实从名字中也能看出，HuBERT的全称是Hidden unit BERT）。
到这里，可能有读者要问了。道理我都懂，可 Z 是怎么得到的呢？下面我就给大家慢慢道来。可以看到上图中的红色区域（Acoustic Unit Discovery System），这一个模块负责的就是给输入的 X’ 进行预分类，得到HuBERT中所谓的Hidden Unit。在预训练之初，HuBERT对于输入w，抽取其39维的MFCC特征，用其进行第一次的k-means聚类（类别数为100），得到初始的伪标签，进行第一轮预训练。然而，仅仅用MFCC这种原始的语音特征进行聚类显然是非常粗糙的。因此，HuBERT采取了一个迭代优化伪标签的策略，使用上一个epoch中的HuBERT模型的某一个中间层的特征作为训练集，重新用k-means训练一组新的Z（类别数为500）。由于HuBERT的中间层的维度至少是768，且包含了更多的上下文关系，利用这些特征得到的聚类更为精准。作者也对使用MFCC的特征和训练一个epoch后的HuBERT中间层特征的聚类效果进行了对比，发现后者的PNMI分数（一种衡量聚类效果的指标）要远远高于前者。
简单介绍完HuBERT的原理，我们来看看在icefall中是如何实现HuBERT的微调的。预训练得到的HuBERT本质上是一个transformer encoder，我们只需要将HuBERT作为端对端ASR模型的encoder，再根据想要的模型框架（Transducer或者CTC）配置相应的decoder即可。本文中我们使用pruned RNNT损失函数对HuBERT进行Transducer框架下的微调，具体的代码实现详见（https://github.com/k2-fsa/icefall/pull/588)。pruned RNNT loss是icefall内设计的一种速度快，省显存的RNNT损失函数，之前小编就已经给大家做过详细介绍啦，有兴趣回顾的读者们请移步多快好省的 RNN-T 训练和Pruned RNN-T 何以又快又好。
在hubert_encoder.py文件中，我们定义了预训练HuBERT模型做为Encoder的HubertEncoder。其定义函数如下：
下面小编带着大家过一遍模型所需要传入的参数。model_dir是本地中储存的预训练好的HuBERT模型（注意，不是微调过后的模型！）所在的路径。output_size是HubertEncoder的输出维度，是一个和HuBERT模型无关的参数。freeze_finetune_updates是训练时的一个超参数，控制HuBERT模型的参数在前多少此更新中保持不变。在HuBERT原作中，作者推荐在微调的初期不更新HuBERT的参数，仅仅更新任务相关模块的参数。等模型预热完毕后，再对HuBERT模型的参数进行更新。接下来三个mask开头的参数mask_prob,mask_chanmnel_prob,mask_channel_length决定了在微调中，多少CNN encoder的输出会被遮盖掉。这样的遮盖类似于端对端ASR模型中最为常见的数据增强做法SpecAug，能够增强模型的鲁棒性。之后两个参数subsample_output是我们自己添加的参数，用来决定是否对HuBERT的输出进行降采样。由于HuBERT中的Transformer的输出频率为50Hz，而一般常见的Encoder模型的输出往往是25Hz。如果subsample_output设置为True，则拼接相邻的两帧输出，并做一次线性变化，将HubertEncoder的输出频率降至25Hz。最后一个参数training判断模型当前是否处在训练状态，如果设置为False（例如推理时），则不会对CNN encoder的输出进行mask操作。
创建好HubertEncoder后，即可指定其作为Transducer模型的encoder。在定义好Transducer内的decoder和joiner，即可以对模型进行微调。在HuBERT论文中，作者使用了一个3分段的学习率调度器（tri-state lr scheduler）对模型进行微调。该调度器在icefall中也有实现（略有改动），代码如下：
TriStateScheduler将学习率的调整分为三段，第一段为热身期（warm up），即学习率从初始学习率init_lr线性增长到warmup_lr；第二段为平稳期，即学习率保持在warmup_lr不变；第三段为下降期（decay），学习率从warmup_lr线性下降到end_lr。每一阶段的步数是由总部数total_steps和phase决定的。例如total_steps=10000，phase=[0.1,0.4,0.5]，则三个阶段的步数分别为10000*0.1=1000，10000*0.4=4000，和10000*0.5=5000。

3. 实验结果

我们使用HuBERT base和HuBERT large在LibriSpeech数据集上进行实验。实验中我们使用了500的BPE作为建模单元，在使用pruned RNNT loss进行微调后，得到了如下的结果：
使用clean-100h进行实验，使用modified_beam_search解码：
参数量test-cleantest-other
HuBERT base96M5.0311.23
HuBERT Large320M3.196.30
使用完整960h进行实验，使用modified_beam_search解码：
参数量test-cleantest-other
HuBERT Base96M2.817.09
Hubert Large320M1.933.93
为了验证icefall中微调后得到的模型的效果，我们与fairseq提供在960h微调好的CTC模型进行对比，同样使用greedy search进行解码，结果如下：
test-cleantest-other
Fairseq HuBERT Large (CTC)2.084.23
Icefall finetuned HuBERT Large

(Transducer, ours)
1.984.10
可以发现，使用pruned RNNT loss对HuBERT进行微调，能够得到比使用CTC作为损失函数更高的识别准确率。这证明icefall中对于HuBERT的微调流程是有效的。

4. 结语

今天给大家简单介绍了HuBERT模型，并且展示了如何在icefall中使用pruned RNNT loss对HuBERT模型进行微调。本文仅是抛砖引玉，欢迎更多的读者参与进来，分享自己对于HuBERT模型的使用经验~

往期文章

Pruned RNN-T 何以又快又好
新一代Kaldi中的RandomCombiner
sherpaC++API之非流式篇
玩转新一代Kaldi流式模型
本文出品：新一代Kaldi-NGK编辑部
撰文：蛋哥的 Xiaoyu Yang
继续阅读
阅读原文