【语音识别】Zipformer

news2024/12/23 13:19:45

Zipformer 是kaldi 团队于2024研发的序列建模模型。相比较于 Conformer、Squeezeformer、E-Branchformer等主流 ASR 模型，Zipformer 具有效果更好、计算更快、更省内存等优点。并在 LibriSpeech、Aishell-1 和 WenetSpeech 等常用数据集上取得了当时最好的 ASR 结果。

一.方法

1. Down sampled encoder structure

2. Zipformer block

3. BiasNorm

4. Swoosh 激活函数

5. ScaledAdam

论文地址：https://arxiv.org/pdf/2310.11230.pdf

项目地址：https://github.com/k2-fsa/icefall/tree/master/egs/librispeech/ASR/zipformer

一.方法

Zipformer的整体框架如下图所示。

不同于 Conformer 只处理固定帧率 25Hz ，Zipformer 采用了1个类似于 U-Net 的结构，在不同帧率上学习时域表征。

首先，Conv-Embed 将输入的 100Hz 的声学特征下采样为 50 Hz 的特征序列；然后，由 6 个连续的 encoder stack 分别在 50Hz、25Hz、12.5Hz、6.25Hz、12.5Hz 和 25Hz 的采样率下进行时域建模。除了第1个 stack 外，其他的 stack 都采用了降采样的结构。在 stack 与 stack 之间，特征序列的采样率保持在 50Hz。不同的 stack 的 embedding 维度不同，中间stack 的 embedding 维度更大。每个 stack 的输出通过截断或者补零的操作，来对齐下1个 stack 的维度。Zipformer 最终输出的维度，取决于 embedding 维度最大的stack。

1. Down sampled encoder structure

•Conv-Embed

使用3个2-D卷积层，其时间×频率步长分别为1×2、2×2和1×2，输出通道分别为8、32和128。随后，利用了一个类似于Nextformer的ConvNeXt层，该层由1个kernel大小为7×7的深度卷积、1个具有384个输出通道的点卷积、1个SwooshL激活函数和1个具有128个输出通道的点卷积组成。在ConvNeXt模块上应用了残差连接。最后，使用1个线性层，后面跟着1个BiasNorm，以调整特征维度，使其与第1个stack相匹配。

•Downsampled stacks

对于降采样的 encoder stack，成对出现的 Downsample 和 Upsample 模块负责将特征长度对称地缩放。当降采样率为 2 时，Downsample 学习2个标量权重用来将相邻的2帧加权求和；Upsample 将每1帧复制为2帧。最后，通过1个 Bypass 模块整合 stack 的输入和输出。

2. Zipformer block

Zipformer block的结构如下图左侧所示。

Zipformer block深度大约是 Conformer block 的2倍。具体地，block 输入先被送到 MHAW 模块计算注意力权重attention weights，attention weights作为NLA 模块和 SA 模块的输入。同时，block 输入也被送到 feed-forward 模块，后接 NLA 模块和2个连续的模块组（SA + convolution + feed-forward）。最后，由1个 BiasNorm 模块对block 输出进行 normalize操作。除了残差连接，每个 Zipformer block 使用2个 Bypass 模型，用于结合 block 输入和中间模块的输出，分别位于 block 的中间和尾部。

•Non-Linear Attention

上图右侧为Non-Linear Attention的结构。利用 MHAW 模块计算好的注意力权重，沿着时间轴汇聚不同帧的向量。具体而言，使用3个 linear 将输入转换为 A、B、C，每个的维度为输入维度的 3/4 倍。模块的输出为 $linear(A\odot attention(\tanh (B)\odot C))$ ，⊙ 表示点乘，attention 表示利用1个注意力头的权重对不同帧汇聚， linear layer 负责恢复特征的维度。

•Bypass

Bypass 模块学习1个逐通道的权重 $c$ ，结合模块输入 $x$ 和模块输出 $y:(1-c)\odot x+c\odot y$ 。在训练早期通过约束 $c$ 的最小值让模块接近 “straight-through” 有助于稳定模型训练。

3. BiasNorm

提出 BiasNorm 模块来替换 LayerNorm：

其中， $b$ 是可学习的逐通道的 bias， $RMS[x-b]$ 是通道的均方根值， $\gamma$ 是1个可学习的标量。

4. Swoosh 激活函数

提出2个新的激活函数用于代替 Swish，分别称为 SwooshR 和 SwooshL。

在 SwooshR 函数中，偏移值 0.313261687 是为了让函数经过原点；在 SwooshL函数中，偏移量 0.035 是经过实验得到的。

如下图所示，SwooshL 近似于 SwooshR 向右偏移得到的。

把 SwooshL 用在 “normally-off” 的模块（feed-forward 和 ConvNeXt）中，把 SwooshR 用在convolution 和 Conv-Embed 中其余的部分。

5. ScaledAdam

提出1个 Adam 优化器的 parameter-scale-invariant 版本，称为 ScaledAdam，可以加快模型收敛。

令 $f(\theta )$ 为我们想要优化的 loss 函数，它对参数 $\theta$ 是可导的。在每个步骤 $t$ ，Adam 计算参数梯度 $g(t)=\bigtriangledown _{\theta }f(\theta _{t-1})$ ，并更新梯度的一阶动量 $m(t)=\beta _{1}m _{t-1} +(1-\beta _{1})g_{t}$ 和二阶动量 $v(t)=\beta _{2}v _{t-1} +(1-\beta _{2})g_{t}^{2}$ ，此处, $\beta _{1}$ , $\beta _{2}$ 表示控制动量更新的系数。Adam 在步骤 t 的参数更新量 $\Delta _{t}$ 为：

$\alpha _{t}$ 通常由外部的 LR schedule 控制， $\frac{\sqrt{1-\beta _{2}^{t}}}{1-\beta _{1}^{t}}$ 为偏置纠正项。

•Scaling update

为了确保不同 scale 的参数的相对变化量 $\frac{\Delta _{t}}{r_{t-1}}$ 一致，在参数更新量中引入参数的 scale，来放缩更新量 $\Delta _{t}$ ：

•Learning parameter scale

从 $r _{t-1}$ 更新到 $r _{t}$ 对参数 $\theta$ 带来的变化为 $\Delta _{t,r}^{'}=(r_{t}-r_{t-1})\odot \theta _{t-1}^{'}$ 。

其中， $\eta$ 是学习率 $\alpha _{t}$ 的缩放参数，值为0.1时有助于稳定训练。

•Eden schedule

Eden schedule的公式如下：

其，t为 step，e为 epoch, $\alpha _{step}$ 和 $\alpha _{epoch}$ 分别控制学习率在哪个 step 和 epoch 开始快速下降,
$linear(\alpha _{start},t _{warmup},t)$ 表示1个线性 warmup，起点为 $\alpha _{start}$ ，经过 $t _{warmup}$ 个 step 变为 1。
$\alpha _{base}$ 表示当没有 warmup 的情况下学习率的最大值。