Zipformer

近日，小米集团新一代 Kaldi 团队关于语音识别声学模型的论文《Zipformer: A faster and better encoder for automatic speech recognition》被 ICLR 2024 接收为 Oral (Top 1.2%)。新型自动语音识别 (ASR) 模型Zipformer诞生之路

• 论文链接：https://arxiv.org/pdf/2310.11230.pdf

• 代码链接：https://github.com/k2-fsa/icefall/tree/master/egs/librispeech/ASR/zipformer

团队介绍

新一代 Kaldi 团队是由 Kaldi 之父、IEEE fellow、小米集团首席语音科学家 Daniel Povey 领衔的团队，专注于开源语音基础引擎研发，从神经网络声学编码器、损失函数、优化器和解码器等各方面重构语音技术链路，旨在提高智能语音任务的准确率和效率。

目前，新一代 Kaldi 项目主要由四个子项目构成：核心算法库 k2、通用语音数据处理工具包 Lhotse、解决方案集合 Icefall 以及服务端引擎 Sherpa，方便开发者轻松训练、部署自己的智能语音模型。

 新一代 kaidi 项目：https://github.com/k2-fsa

论文解读

摘要

Zipformer[1] 作为一个新型的自动语音识别 (ASR) 模型，相比较于 Conformer[2]、Squeezeformer[3]、E-Branchformer[4] 等主流 ASR 模型，Zipformer 具有效果更好、计算更快、更省内存等优点。Zipformer 在 LibriSpeech、Aishell-1 和 WenetSpeech 等常用的 ASR 数据集上都取得了当前最好的实验结果。

Zipformer 的具体创新点，主要包括：

• Downsampled encoder structure，降采样到不同的帧率，学习不同粒度的时域信息

• Zipformer block，更深的 block 结构，通过复用注意力权重提高效率

• BiasNorm，允许保留一定的长度信息

• Swoosh 激活函数，比 Swish 效果更好

• ScaledAdam 优化器，根据参数大小放缩更新量，保持不同参数相对变化一致，并显式学习参数大小，比 Adam 收敛更快、效果更好

• Balancer 和 Whitener，限制激活值，稳定训练

方法

1. Downsampled encoder structure

图 1 展示了 Zipformer 总体框架图，由一个 Conv-Embed 模块和多个 encoder stack 组成。不同于 Conformer 只在一个固定的帧率 25Hz 操作，Zipformer 采用了一个类似于 U-Net 的结构，在不同帧率上学习不同时间分辨率的时域表征。

图1：Zipformer 总体框架

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

对于降采样的 encoder stack，成对出现的 Downsample 和 Upsample 模块负责将特征长度对称地放缩。我们采用几乎最简单的方法实现 Downsample 和 Upsample 模块。例如，当降采样率为 2 时，Downsample 学习两个标量权重，用来将相邻的两帧加权求和了；Upsample 则只是简单地将每一帧复制为两帧。最后，通过一个 Bypass 模块，以一种可学习的方式结合 stack 的输入和输出。

2. Zipformer block

Conformer block 由四个模块组成：feed-forward、Multi-Head Self-Attention (MHSA)、convolution、feed-forward。MHSA 模块通过两个步骤学习全局时域信息：基于内积计算注意力权重，以及利用算好的注意力权重汇聚不同帧的信息。然而，MHSA 模块通常占据了大量的计算，因为以上两步操作的计算复杂度都是平方级别于序列长度的。因此，我们将 MHSA 模块根据这两个步骤分解为两个独立的模块：Multi-Head Attention Weight (MHAW) 和 Self-Attention (SA)。 这样一来，我们可以通过在一个 block 里面使用一个 MHAW 模块和两个 SA 模块，以高效的方式实现两次注意力建模。此外，我们还提出了一个新的模块 Non-Linear Attention (NLA) ，充分利用已经算好的注意力权重，进行全局时域信息学习。

图2：Zipformer block 结构图

图 2 展示了 Zipformer block 的结构图，其深度大约是 Conformer block 的两倍。核心的思想是通过复用注意力权重来节省计算和内存。 具体而言，block 输入先被送到 MHAW 模块计算注意力权重，并分享给 NLA 模块和两个 SA 模块使用。同时，block 输入也被送到 feed-forward 模块，后面接着 NLA 模块。接着是两个连续的模块组，每组包含 SA、convolution 和 feed-forward。最后，由一个 BiasNorm 模块来将 block 输出进行 normalize。除了普通的加法残差连接，每个 Zipformer block 还使用了两个 Bypass 模型，用于结合 block 输入和中间模块的输出，分别位于 block 的中间和尾部。

值得注意的是，我们并没有像常规的 Transformer 模型一样，对每个模块都使用 normalization layer 去周期性地调整激活值的范围，这得益于我们使用的 ScaledAdam 优化器可以为各个模型自动学习参数的 scale。

Non-Linear Attention

图 3：Non-Linear Attention 模块

Bypass

4. Swoosh 激活函数

Conformer 采用的激活函数为 Swish[6]，其公式为：

在 SwooshR 函数中，偏移值 0.313261687 是为了让函数经过原点；在 SwooshL函数中，偏移量 0.035 是经过实验调过的，比让它精确经过原点的结果略好。  whaosoft aiot http://143ai.com

如图 4 所示，SwooshL 近似于 SwooshR 向右偏移得到的。“L” 和 “R” 表示两个与 x 轴的交点中哪个靠近或经过原点。类似于 Swish，SwooshR 和 SwooshL 都是有下确界的并且非单调的。相比较于 Swish，最大的区别在于 SwooshR 和 Swoosh 对于负数部分有一个斜率，这个可以避免输入一直是负数以及避免 Adam-type 的更新量分母（二阶梯度动量）太小。 当将 SwooshR 用在 Zipformer 各个模块中时，我们发现，那些带残差的模块， 例如 feed-forward 和 ConvNeXt，倾向于在激活函数前面的线性层中学习一个绝对值较大的负数 bias，以学习一种 “normally-off” 的行为。 因此，我们把 SwooshL 函数用在这些 “normally-off” 的模块中，把 SwooshR 用在其他的模块：convolution 和 Conv-Embed 剩下的部分。

图4：激活函数 Swish，SwooshR 和 SwooshL

5. ScaledAdam

我们提出一个 Adam 优化器[7] 的 parameter-scale-invariant 版本，称为 ScaledAdam，它可以加快模型收敛。一方面，ScaledAdam 根据参数 scale 放缩参数更新量，来确保不同 scale 的参数相对变化一致；另一方面，ScaledAdam 显式学习参数的 scale，这相当于给了一个额外的放缩参数 scale 的梯度。

Learning parameter scale

Eden schedule

Eden schedule 的公式如下：

Efficient implementation

实验

1. 实验设置

Architecture variants

我们构建了三个不同参数规模的 Zipformer 模型：small (Zipformer-S)， medium (Zipformer-M)，large (Zipformer-L)。对于 Zipformer 的 6 个 stack，注意力头的数量为 {4,4,4,8,4,4}，卷积核大小为 {31,31,15,15,15,31}。对于每个注意力头，query/key 维度为 32，value 维度为 12。我们通过调节 encoder embedding dim，层的数量，feed-forward hidden dim 来得到不同参数规模的模型：

表 1：不同规模 Zipformer 的参数配置

数据集

我们在三个常用的数据集上进行实验：1）Librispeech[8]，1000 小时英文数据；2）Aishell-1[9]，170 小时中文；3）WenetSpeech[10]，10000+ 小时中文数据。

实现细节

我们通过 Speed perturb 对数据进行三倍增广，使用 Pruned transducer[11] 作为 loss 训练模型，解码方法为 modified-beam-search[12]（每帧最多吐一个字，beam size=4）。

默认情况下，我们所有的 Zipformer 模型是在 32GB NVIDIA Tesla V100 GPU 上训练。对于 LibriSpeech 数据集，Zipformer-M 和 Zipformer-L 在 4 个 GPU 上训练了 50 epoch，Zipformer-S 在 2 个 GPU 上训练了 50 个 epoch；对于 Aishell-1 数据集，所有 Zipformer 模型都在 2 个 GPU 上训练了 56 epoch；对于 WenetSpeech 数据集，所有 Zipformer 模型都在 4 个 GPU 上训练了 14 epoch。

2. 与 SOTA 模型比较

LibriSpeech

表 2 展示了 Zipformer 和其他 SOTA 模型在 LibriSpeech 数据集上的结果。对于 Conformer，我们还列出了我们复现的结果以及其他框架复现的结果。值得注意的是，这些结果和 Conformer 原文仍然存在一定的差距。Zipformer-S 取得了比所有的 Squeezeformer 模型更低的 WER，而参数量和 FLOPs 都更少。Zipformer-L的性能显著超过 Squeezeformer-L，Branchformer 和 我们复现的 Conformer，而 FLOPs 却节省了 50% 以上。值得注意的是，当我们在 8 个 80G NVIDIA Tesla A100 GPU 上训练 170 epoch，Zipformer-L 取得了 2.00%/4.38% 的 WER，这是我们了解到的迄今为止第一个和 Conformer 原文结果相当的模型。

表2：不同模型在 LibriSpeech 数据集的比较

我们还比较了 Zipformer 和其他 SOTA 模型的计算效率和内存使用。图 5 展示了不同 encoder 在单个 NVIDIA Tesla V100 GPU 上推理 30 秒长的语音 batch 所需的平均计算时间和峰值内存使用量，batch size 设置为 30，确保所有的模型都不会 OOM。总的来说，与其他的 SOTA 模型比较，Zipformer 在性能和效率上取得了明显更好的 trade-off。尤其是 Zipformer-L，计算速度和内存使用显著优于其他类似参数规模的模型。

此外，我们在论文附录中也展示了 Zipformer 在 CTC 和 CTC/AED 系统中的性能，同样超过了 SOTA 模型。CTC/AED 的代码在 https://github.com/k2-fsa/icefall/pull/1389。

图 5：不同模型的计算速度和内存使用比较

Aishell-1

表 3 展示了不同模型在 Aishell-1 数据集的结果。相比较于 ESPnet 框架[13] 实现的 Conformer，Zipformer-S 性能更好，参数更少。增大参数规模后，Zipformer-M 和 Zipformer-L 都超过了其他所有的模型。

表 3：不同模型在 Aishell-1 数据集的比较

WenetSpeech

表 4 展示了不同模型在 WenetSpeech 数据集的结果。Zipformer-M 和 Zipformer-L 都在 Test-Net 和 Test-Meeting 测试集上超过了其他所有的模型。Zipformer-S 的效果超过了 ESPnet[13] 和 Wenet[14] 实现的 Conformer，参数量却只有它们的 1/3。

表 4：不同模型在 WenetSpeech 数据集的比较

目前，除了论文中展示的 LibriSpeech、Aishell-1 和 WenetSpeech 数据集外，我们的实验表明， Zipformer 在其它较大规模的 ASR 数据集上同样取得了新的 SOTA 结果。例如在 10000 h 的英文数据集 GigaSpeech[15] 上，不使用外部语言模型时，在 dev/test 测试集上，66M Zipformer-M 的 WER 为 10.25/10.38，288M Zipformer 的 WER 为 10.07/10.19。

3. 消融实验

我们在 LibriSpeech 数据集上进行了一系列消融实验，验证每一个模块的有效性，实验结果如表 5 所示。

表 5：Zipformer 消融实验

Encoder structure

我们移除了 Zipformer 的 Downsampled encoder structure，类似于 Conformer 在 Conv-Embed 中使用 4 倍降采样，得到一个 12 层的模型，每层的 embedding dim 为 512。该模型在两个测试集上的 WER 都有所上升，这表明 Zipformer 中采用的 Downsampled encoder structure 并不会带来信息损失，反而以更少的参数达到更好的性能。

Block structure

由于每个 Zipfomer block 含有两倍于 Conformer block 的模块数量，我们将每个 Zipformer block 替换为两个 Conformer block，这导致了在 test-other 上的 WER 上升了 0.16%，并且带来更多的参数量，这体现了 Zipformer block 的结构优势。移除 NLA 或者 Bypass 模块都导致了性能下降。对于移除了 NLA 的模型，当我们移除注意力共享机制，这并没有带来性能提升，反而会带来更多的参数和计算量。我们认为在 Zipformer block 中两个注意力模块学习到的注意力权重具有高度一致性，共享注意力权重并不会有损模型性能。

Normalization layer

将 BiasNorm 替换为 LayerNorm 导致在 test-clean 和 test-other 两个测试集上 WER 分别上升了 0.08% 和 0.18%，这表明了 BiasNorm 相对于 LayerNorm 的优势，可以对输入向量保留一定程度的长度信息。

Activation function

当给 Zipformer 所有的模块都是用 SwooshR 激活函数的时候，test-clean 和 test-other 两个测试集上 WER 分别上升了 0.11% 和 0.42%，这表明给那些学习 “normally-off” 行为的模块使用 SwooshL 激活函数的优势。给所有的模块使用 Swish 激活函数导致了更严重的性能损失，这体现了 SwooshR 相对于 Swish 的优势。

Optimizer

图 6：ScaledAdam 与 Adam 的比较

Activation constraints

如表 6 所示，我们将 Balancer 移除掉后并不会带来明显的性能变化，但是没有对激活值的范围作限制会增大模型不收敛的风险，尤其是在使用混合精度训练的时候。移除掉 Whitener 导致了在 test-clean 和 test-other 两个测试集上分别下降了 0.04% 和 0.24%，这表明通过限制激活值的协方差矩阵特征值尽可能相同，有助于让提升模型性能。

表 6：Activation constraints 消融实验

结语

目前，Zipformer 已在小米产线数据上充分验证了其优越性能，有效提升识别准确率，降低服务器成本。Zipformer 相关技术，如 ScaledAdam 等，已被用于小米大模型训练。另外，我们的初步实验表明，Zipformer 在视觉模型上同样展示了有效性。