小米新一代Kaldi团队论文解读：新型自动语音识别 (ASR) 模型Zipformer诞生之路｜ICLR 2024 Oral

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

• 论文链接：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，方便开发者轻松训练、部署自己的智能语音模型。

论文解读

摘要

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 的结构，在不同帧率上学习不同时间分辨率的时域表征。

首先，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 的结构图，其深度大约是 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 展示了 NLA 模块的结构。类似于 SA 模块，它利用 MHAW 模块计算好的注意力权重，沿着时间轴汇聚不同帧的向量。 具体而言，它使用三个 linear 将输入转换为 A、B、C，每个的维度为输入维度的 3/4 倍。模块的输出为 ，⊙ 表示点乘， 表示利用一个注意力头的权重对不同帧汇聚， 负责恢复特征的维度。

                                 ^{图 3：Non-Linear Attention 模块}

Bypass 模块学习一个逐通道的权重，结合模块输入和模块输出 。我们发现，在训练早期通过约束 的最小值让模块接近 “straight-through” 有助于稳定模型训练。

3. BiasNorm

Conformer 使用 LayerNorm[5] 来 normalize 激活值，给定维的向量 。LayerNorm 的公式为：

LayerNorm 先计算均值和方差，用于向量标准化，将向量长度调整为 。然后，利用逐通道的放缩因子和偏置进行元素变换，这有助于调整不同模块对整个模型的相对贡献。

然而，我们观察到使用 LayerNorm 的 Conformer 模型存在着两种失败的情况：1）有时候会将某个特征维度设置得非常大，例如 50 左右，我们认为这是模型在抵制 LayerNorm 完全消除长度的机制，这个非常大的数可以在 normalize 的过程中保留其他维度的一部分长度信息。2）有些模块（例如 feed-forward 和 convolution）的输出值非常小，例如 1e-6。我们认为在模型开始训练的时候，还没有学到有用信息的模块被防缩因子通过接近 0 关闭了。如果放缩因子在 0 左右震荡，反向传播的梯度也会随之翻转，这样一来，模块很难学到有用的信息，因为这是一个难以跳出的局部鞍点。

其中，是可学习的逐通道的 bias，是一个可学习的标量。首先，我们去除了减均值的操作，因为它没有必要，除非它接着一个非线性变换。充当一个非常大的数，用于在 normalize 的过程中保留向量的一部分长度信息， 这样一来，模型就不需要牺牲一个额外的维度来做这个事情。这或许有助于模型量化，因为它可以减少离群点的出现。由于一直是正数，避免了出现梯度方向翻转导致的某些模块无法学习的问题。

4. Swoosh 激活函数

5. ScaledAdam

令为我们想要优化的 loss 函数，它对参数是可导的。在每个步骤，Adam 计算参数梯度，并更新梯度的一阶动量和二阶动量，此处, 表示控制动量更新的系数。Adam 在步骤 t 的参数更新量为：

式中，通常由外部的 LR schedule 控制，为偏置纠正项。尽管 Adam 对梯度 scale 是 invariant 的，但是我们认为它仍然存在两个问题：1）更新量并没有考虑参数的 scale（标记为），对于参数的相对更新量而言，Adam 可能会导致对 scale 小的参数学习太快，或者对 scale 大的参数学习太慢。2）我们很难直接学习参数的 scale，因为参数 scale 的大小变化方向是高维度的梯度向量中一个特别具体的方向。 尤其是 scale 变小的方向更加难学，因为在优化的过程中，梯度会引入噪声，参数的 scale 会倾向于不断增大。

为了确保不同 scale 的参数的相对变化量一致，我们在参数更新量中引入参数的 scale，来放缩更新量：

我们计算作为参数的 scale。由于 ScaledAdam 比 Adam 更不容易发散，我们使用一个不需要很长 warm-up 的 LR schedule，称为 Eden；我们使用明显更大的学习率，因为参数的 RMS 值通常小于 1。

为了显式学习参数的 scale，我们在将它当作一个真的存在的参数一样学习，仿佛我们将每个参数分解为，并且我们是对参数 scale 和内部参数进行梯度下降。 值得注意的是，在具体实现中，我们并没有将每个参数进行分解，只是增加了一个额外的更新参数 scale 的梯度。

令为参数 scale 的梯度，存在。由于 Adam 对梯度的 scale 几乎是 invariant 的，我们可以计算 。按照 Adam 算法，我们维护参数 scale 梯度的一阶动量和二阶动量。将参数 scale 从更新到对参数带来的变化为。同样地，我们放缩参数 scale 对应的更新量：

式中， 用于放缩学习率，我们发现设置为 0.1 有助于稳定训练。此时，参数的更新量由变为，这等价于额外引入一个放缩参数 scale 的梯度。 这个改动有助于我们简化模型结构，我们可以去掉大部分的 normalization 层，因此每个模块可以更容易得学习参数 scale ，来将激活值调整到一个合适的范围。

式中，为 step，为 epoch，和分别控制学习率在哪个 step 和 epoch 开始快速下降， 表示一个线性 warmup，起点为，经过个 step 变为 1。表示当没有 warmup 的情况下学习率的最大值。让 Eden 同时依赖于 step 和 epoch 两个变量，是为了让模型的更新程度在经过一定的训练数据量（e.g., 1h）时，几乎不受 batch size 影响。 Eden 公式中，epoch 也可以替换为其他合适的变量，如经过多少小时的数据。

为了加快 ScaledAdam 计算，我们将参数根据 shape 分组，按照 batch 进行参数更新， 值得注意的是这并不影响结果。Scaleadam 需要的内存使用量和 Adam 差不多，只需要额外的内存来存储参数 scale 对应梯度的一阶动量和二阶动量和 。

为了确保训练的一致性以及避免训练出性能差的模型，我们提出 Balancer 和 Whitener，用于约束模型的激活值。 Balancer 和 Whitener 以一种省内存的方式实现：在前向过程中，相当于是一个 no-op；在反向过程中，计算一个对激活值施加限制的损失函数的梯度，加到原本的激活值梯度上：。 Balancer 和 Whitener 的应用位置没有遵循一个明确的规则，我们一般是在模型表现不好的时候，通过分析模型哪个地方出现问题，再对应地使用 Balancer 和 Whitener 去修复模型。

Balancer 通过对激活值施加限制：最小和最大平均绝对值， 分别标记为和；最小和最大正数比例， 分别标记为和。由于正数比例是不可导的，我们将限制转化为 standard-deviation-normalized mean ，得到 和 。同时，我们将平均绝对值转化为 RMS：，得到和。具体而言，对于激活值，限制函数定义为：

式中，，和为每个通道的统计量。

Whitener 

Whitener 模块旨在通过限制协方差矩阵的特征值尽可能相同，来鼓励模块学习更有信息量的输出分布。 具体而言，对于特征 ，我们计算协方差矩阵， 为各个通道的均值。Whitener 定义的限制函数为：

式中，为协方差矩阵的特征值。

实验

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 来得到不同参数规模的模型：

2. 与 SOTA 模型比较

LibriSpeech 

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

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

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. 消融实验

当我们使用 Adam 来训练 ScaledAdam 的时候，我们必须给每个模块配一个 BiasNorm 来防止模型不收敛，因为 Adam 无法像 ScaledAdam 一样很好地学习参数 scale 来放缩激活值的大小。我们给两个优化器都尝试了不同的学习率：ScaledAdam（0.025, 0.035, 0.045, 0.055），Adam（2.5, 5.0, 7.5, 10.0）。我们给 Adam 使用与 Conformer 原文[2]一样的 LR schedule：。图 6 分别展示了使用 ScaledAdam 和 Adam 的模型在不同 epoch 时的平均 WER，以及对应的学习率，我们将它们最好的结果分别展示在表 5 中。与 Adam 训练的模型相比较，使用 ScaledAdam 训练的模型在 test-clean 和 test-other 两个测试集上的 WER 分别下降了 0.17% 和 0.72%，同时 ScaledAdam 收敛更快、效果更好。

结语

^参考文献

^{[1] Zipformer: A faster and better encoder for automatic speech recognition (https://arxiv.org/pdf/2310.11230)}

^{[2] Conformer: Convolution-augmented Transformer for Speech Recognition (https://arxiv.org/abs/2005.08100)}

^{[3] Squeezeformer: An Efficient Transformer for Automatic Speech Recognition (https://arxiv.org/abs/2206.00888)}

^{[4] E-Branchformer: Branchformer with Enhanced merging for speech recognition (https://arxiv.org/abs/2210.00077)}

^{[5] Layer Normalization (https://arxiv.org/abs/1607.06450)}

^{[6] Swish: a Self-Gated Activation Function (https://arxiv.org/abs/1710.05941v1)}

^{[7] Adam: A Method for Stochastic Optimization (https://arxiv.org/abs/1412.6980)}

^{[8] LibriSpeech: An ASR corpus based on public domain audio books (https://danielpovey.com/files/2015_icassp_librispeech.pdf)}

^{[9] Aishell-1: An open-source mandarin speech corpus and a speech recognition baseline (https://arxiv.org/abs/1709.05522)}

^{[10] WenetSpeech: A 10000+ Hours Multi-domain Mandarin Corpus for Speech Recognition (https://arxiv.org/abs/2110.03370)}

^{[11] Pruned RNN-T for fast, memory-efficient ASR training (https://arxiv.org/abs/2206.13236)}

^{[12] Fast and parallel decoding for transducer (https://arxiv.org/abs/2211.00484)}

^{[13] ESPnet: https://github.com/espnet/espnet}

^{[14] Wenet: https://github.com/wenet-e2e/wenet}

^{[15] GigaSpeech: An Evolving, Multi-domain ASR Corpus with 10,000 Hours of Transcribed Audio (https://arxiv.org/abs/2106.06909)}