本文将介绍新一代 Kaldi 团队在 Conformer 上的改进工作。
更多内容建议参考:
Reworked Conformer 模型代码[1](https://github.com/k2-fsa/icefall/blob/master/egs/librispeech/ASR/pruned_transducer_stateless2/)

简 介

Conformer[2] 是谷歌团队在 2020 年提出的用于语音识别的声学模型。当前,我们主要使用 Conformer 作为声学模型,结合 Pruned RNN-T[3] 构建语音识别模型,已取得了不错的效果。我们使用的 Conformer 模型包含了 12 个 encoder layer,每个 encoder layer 包含了四个模块:Feed-forward + Self-attention + Convolution + Feed-forward,每个模块的输入处以及 encoder layer 的最后输出处各有一个 LayerNorm。
通过使用模型分析工具 diagnostics[4] 对 Conformer 进行分析,我们发现了模型存在的一些问题,并针对问题作出改进,得到 Reworked Conformer。 
模型分析工具 diagnostics 实现于 icefall[5] 中,可用于查看模型的参数、激活值及梯度的多种统计量,方便用户分析模型是否训练正常。

Reworked Conformer 改进之处

1. 将 LayerNorm 替换为 BasicNorm

为什么不使用 BatchNorm

对于大小为 的输入特征 , 其中 为 minibatch 大小, 为长度, 为特征通道数,BatchNorm 和 LayerNorm 的计算过程均可表达为如下公式:
公式中主要包括两步操作:
  • 计算均值 和方差 ,进行标准化。
  • 利用可学习的参数 和 ,调整值的范围。
BatchNorm 和 LayerNorm的主要区别在于:
  • BatchNorm 的统计维度为 ,在训练过程中累计统计量(均值、方差),测试时使用训练中累计的统计量;
  • LayerNorm 的统计维度为 ,无论训练或测试,都要重新计算统计量。
BatchNorm 假设了 minibatch 中不同的句子(样本)服从独立同分布(iid)。然而,该假设在ASR模型训练过程中,可能会遇到不成立的情况,例如:
  • 按照句子长度划分 minibatch;
  • 使用多个数据集训练,在不同 minibatch 使用不同数据集。
因此,我们首先将 Conformer 中所有的 BatchNorm 替换为 LayerNorm。

为什么需要 LayerNorm

LayerNorm 的作用主要包括:
  • 规范化值的范围,避免在训练中因为值太大或太小而导致 loss 出现 nan 或者 inf ;对于任意长度的向量,LayerNorm 会将其长度规范化为 ,其中 为特征通道个数。
  • 利用可学习的参数 ,调整对应模块对整体模型的贡献。

使用 LayerNorm 时存在的问题

  • Module Death
在开始训练阶段,有些模块还没学习到有用的信息,LayerNorm 中可学习的参数 会企图去关闭掉那些没用的模块。此时,参数 会在 0 左右震荡,模块的梯度符号也随之翻转,导致该模块难以学习。
  • Large constant channel
观察每一个 channel,发现总是存在一个 channel 会输出一个绝对值很大且恒定的值:-51+-1。出现这个现象的原因可能是,模型企图抵抗 LayerNorm 去除向量长度的机制,这个恒定值类似于一个非常大的 ,以达到保留(除了该channel之外的)向量长度的目的。
以下图为例,向量 和向量 经过长度规范化后,长度关系丢失;向量 和向量 的最后一个 channel是一个绝对值很大的恒定值 -51,对于除了该 channel 外的所有 channel,两者之间的长度关系 在长度规范化的同时也得到保留。
BasicNorm
针对上述问题,对 LayerNorm 作出修改,得到 BasicNorm
  • 首先,去掉了分子中移除向量均值的操作,因为该操作是一个线性变换,其前面接的是线性层,没有必要进行连续的两次线性变换。
  • 在分母中,引入一个可学习的参数 ,充当上述的 large constant channel 的角色,目的是在作长度规范化的同时,保留向量的长度信息。此时,输入的向量长度越大,输出向量的长度就越大。
  • 初始值为 0.25,学习的是该参数的 log 值,用的时候取 exp,确保是正数。
下图展示了每个 encoder layer 原来的 (左边)和 reworked 的 (右边)结构差异:

Learnable scales

现在 BasicNorm 中没有可学习的放缩因子参数 ,为了保留其权衡不同模块贡献的功能,在每个模块中引入可学习的参数 scales
self.weight -> self.weight * self.weight_scale.exp()

self.bias -> self.bias * self.bias_scale.exp()

由于 scales 是在 log 空间学习,使用时取 exp,为正数,避免了上述由于在 0 左右震荡导致的对应模块参数梯度翻转的问题。

2. ActivationBalancer调整激活值范围

激活值异常问题

Conformer中使用非线性函数 作为激活函数:
经观察发现,主要存在两种导致激活值异常的情况:
  • Dead neuron
在 Feed-forward(LayerNorm + Linear + Swish + Linear)模块中,对于 之前的 Linear 输出,存在很大比例的负数,有时候达到 50%。然而,负数经过 会接近 0,这造成了大量的参数浪费。
  • Wrong dynamic range
另一种导致激活值异常情况是, 太大或者太小。若 太小,如 1e-5, 会接近线性函数;若 太大,如 100, 会退化为 。

ActivationBalancer

针对上述激活值异常问题,我们设计了 ActivationBalancer 模块,具体而言:
  • 在前向计算的过程中,统计激活值的范围,包括正数的比例以及绝对值的大小。
  • 在反向求梯度的过程中,根据前向统计结果,对应地放缩梯度大小。
例如,当统计发现激活值太小的时候:
  • 若梯度为负,激活值会朝着数轴的正方向走,这是我们所希望的,此时应当放大梯度;
  • 若梯度为正,激活值会朝着数轴的负方向走,这不是我们所希望的,此时应当缩小梯度。
其他情况可以此类推,不再赘述。

3. 效果更好的DoubleSwish

经实验发现,连续应用两次 ,相比较于单次,会有轻微的结果提升。因此,我们使用一个近似于 的函数 (如下图所示),代替 :
的优点在于,不仅能收获轻微的性能提升,在训练过程中比更省内存,原因是该函数的梯度为:
我们只需要缓存 和 ,不需要缓存 。

4. 固定参数范数,让学习率控制参数相对变化

随着训练进行,模型参数范数 RMS(Root-Mean-Square)会随之变大,此时学习率无法直接控制模型参数的相对变化大小。尽管,在参数范数 RMS 增大时,LayerNorm 或BatchNorm 中可学习的参数 可以通过缩小来控制值的范围,但是 太小的时候会导致上述的训练不稳定的问题。
因此,我们在优化器 AdamW[6] 的基础上修改得到优化器 Eve,其核心思想是对于每个非标量参数,当其范数 RMS 超过 0.1 时,对其进行 weight decay,缩小参数范数。
固定参数范数的操作并不会影响模型的性能,原因是我们在每个模块中引入了可学习的参数 scales。

5. 使用Model-level Warmup 机制,加快收敛

在传统的Transformer训练中,通常采用先升后降的学习率 warmup 机制,这是一种在训练深层模型中常用的策略。如下图所示,在前面 80,000 个 iteration 线性地增加学习率,然后再按照 减少学习率。
我们提出 model-level warmup 机制,代替上述的学习率 warmup 机制。在前 3,000 个 iteration 中,逐步将 warmup 值从 0.1 增加到 1.0,在每一个 encoder layer 中,返回值为。该机制相当于在早期的训练阶段,减小模型的有效深度,使得训练更加容易。
此时,我们不再需要采用先升后降的学习率 warmup 机制。我们在开始的时候使用较大的学习率,然后逐步降低学习率,如下图所示:
我们采用的学习率更新策略,兼顾了当前的 epoch 计数 和 iteration 计数 :
该策略的优势在于,由于同时考虑了 epoch 计数,不再需要因为 minibatch 大小改变而调整学习率更新策略。

实验结果

更快的收敛速度和训练速度

如下图所示,相比较于原来的模型(橙色),Reworked Conformer(蓝色)在训练开始阶段的收敛速度显著提升,且能收敛到更低的 loss 值。
(小编觉得有点东西!)
由于 Reworked Conformer 中移除了大部分 LayerNorm ,对于每个 minibatch,训练速度大约加快了 10%。

更低的WER

下面的表格比较了原来的 Conformer 和 Reworked Conformer 的 WER。两者都采用了 Pruned RNN-T 损失函数,在 full-librispeech 数据集上训练了 26 个 epoch (avg-8),使用 greedy search 解码方法。可以看出,Reworked Conformer 在 两个测试集上都取得了更低的 WER。
Method test-cleantest-other
Conformer2.886.96
Reworked Conformer2.626.37

总 结

本文介绍了新一代 Kaldi 改进后的 Corformer 模型 —— Reworked Conformer,该模型目前已被用于 icefall 的多个 recipe 中,经过大量实验验证了其优越性、稳定性和可靠性。其中涉及到的改进思想,不仅限于 Conformer,也可泛化到其他模型,欢迎大家尝试。

往期文章

本文出品:新一代Kaldi-NGK编辑部
撰文:蛋哥的Zengwei Yao,王全东

参考资料

[1]
Reworked Conformer模型代码: https://github.com/k2-fsa/icefall/blob/master/egs/librispeech/ASR/pruned_transducer_stateless2/
[2]
Conformer: https://arxiv.org/abs/2005.08100
[3]
Pruned RNN-T: https://mp.weixin.qq.com/s/bgJHwHp0PyFy0pWGVWvv0w
[4]
diagnostics: https://github.com/k2-fsa/icefall/blob/master/icefall/diagnostics.py
[5]
icefall: https://github.com/k2-fsa/icefall/
[6]
AdamW: https://arxiv.org/abs/1711.05101
继续阅读
阅读原文