USM-Lite：使用通用语音模型进行语音识别的量化和稀疏感知微调

摘要

随着大规模通用语音模型 (USM) 的最新发展，端到端自动语音识别 (ASR) 模型的质量得到了革命性的提升。 然而，由于巨大的内存使用和计算成本，部署这些大规模 USM 的成本极其昂贵。 因此，模型压缩是在现实场景中在预算范围内适应基于 USM 的 ASR 的一个重要研究课题。 在本研究中，我们提出了一种用于 ASR 的 USM 微调方法，在模型权重上采用低位量化和 $N$:$M$ 结构化稀疏感知范式，降低了模型复杂性从参数精度和矩阵拓扑的角度来看。 我们使用 20 亿参数 USM 在大规模语音搜索数据集上进行了广泛的实验，以评估我们提出的方法。 一系列消融研究验证了高达 int4 量化和 2:4 稀疏性的有效性。 然而，在 int2 量化和 1:4 稀疏度等极端设置下，单一压缩技术无法很好地恢复性能。 相比之下，我们提出的方法可以将模型的大小压缩到 9.4%，而相对字错误率 (WER) 回归的代价仅为 7.3%。 我们还对结果进行了深入分析，并对局限性和潜在解决方案进行了讨论，这对未来的研究很有价值。

索引词— 语音识别、模型量化、模型剪枝、稀疏性、通用语音模型

1简介

端到端自动语音识别[1,2,3,4,5]是一种将音频转录为文本的技术，已广泛集成到现代用户交互式人工智能服务和设备中（例如，语音识别）。 、语音搜索、语音助手等）。 在过去几年中，ASR 模型在不同的测试条件下都取得了质量和延迟的改善[6,7,8,9]。 与此同时，事实证明，从质量和生产的角度来看，端到端 ASR 都可以从自监督学习 (SSL) 语音表示[10,11,12,13]中获得极大的好处。 对这些自监督语音模型进行微调可显着提高 ASR 质量。 更重要的是，它们提供了合适的初始化并降低了所有下游任务的训练成本。

最近，随着大容量硬件的迅速出现和大规模数据集的可用性，SSL 语音模型呈现出越来越大的趋势[14,15,16]。 这些模型扩展了传统的 SSL 语音模型，以捕获多域和多语言分布。 凭借这种能力，它们可以作为大多数语音处理任务的通用基础模型。 然而，由于需要大量的内存和计算单元，这些模型的巨大规模（数十亿个参数）使得它们的部署成本极其昂贵。 因此，高效的微调和模型压缩算法已成为前所未有的重要研究课题。

从之前的研究来看，我们通过稀疏网络剪枝[17,18,19,20,21]和模型量化[22,23,24在端到端ASR压缩方面取得了成功, 25] 然而，压缩这些大规模通用语音模型可能会给常规端到端模型带来新的挑战。 例如，USM的模型尺寸很大，因此需要更高的压缩比才能满足部署的效率要求。 更重要的是，大多数当前的压缩方法在高压缩比（例如，>75% 的稀疏度、2 位量化）下会出现相当大的质量回归，这可能会导致部署到生产环境中时的用户体验较差。

这促使我们同时从不同角度研究压缩模型的有效性：量化从参数精度上降低模型复杂度，而稀疏性则侧重于矩阵拓扑。 因此，我们提出了一种用于 ASR 的 USM 微调方法，具有低位量化和 $N$:$M$ 模型权重的结构化稀疏[26]感知范例，其中两种技术都是硬件友好的，并且受到现代 GPU 和 TPU 的支持。 在前向传播期间，我们首先对权重矩阵执行幅度剪枝 $N$:$M$ 稀疏性，然后，我们对非零权重执行每通道量化。 在反向传播时，我们采用直通估计器（STE）[27]来绕过量化舍入函数。

我们对来自[15]的具有20亿参数CTC USM主干的大规模语音搜索数据集进行了广泛的实验，以评估所提出的方法。 通过我们提出的 4 位量化方法和 2:4 稀疏感知训练方案，模型大小可以减小到 float32 大小的 9.4%，而相对 WER 回归的成本仅为 7.3%。 此外，我们通过消融研究系统地基准化和分析量化和稀疏性的有效性，并与相关研究讨论结果。 这些基准和讨论为未来的研究和生产提供了有用的信息。

2相关工作

ASR 的自我监督学习。施耐德等人 [10]提出了第一个关于SSL语音模型的范例研究，该模型通过对比损失直接从原始波形中学习语音表示。 后续研究通过高级学习范式[11,12,13]以及数据和模型大小扩展[14,15,16]改进了SSL语音模型。 从 SSL 预训练编码器进行初始化显着提高了最先进的 ASR 性能，尤其是在低监督数据设置下。

通过量化和稀疏性进行 ASR 模型压缩。 在预算范围内拟合 ASR 模型的标准设计方法是将模型量化或网络剪枝应用于大型模型。 最近的量化研究[22,23,24,25]表明，可以将 ASR 模型量化为 4 位甚至 2 位，而性能损失却很小。 同样，在网络剪枝方面，通过基于迭代幅度剪枝的各种算法，非结构化和结构化稀疏[17,18,19,20,21]模式在高稀疏度水平上都取得了合理的性能。

与之前工作的关系。 之前有几项研究通过稀疏性、知识蒸馏、注意力重用或其组合来研究 SSL 语音模型压缩[28,20,29,30]。 我们提出的研究在几个方面与他们不同。 首先，我们探索稀疏性和量化压缩的结合，这在以前的研究中尚未研究过。 与之前的研究相比，量化中的本机操作在训练和推理时间方面都提供了显着的加速。 更重要的是，所有这些先前的研究都表明在高压缩比下质量会显着下降。 相比之下，我们的模型在原始模型大小的 9.4% 下只有 7.3% 的相对 WER 回归。 此外，我们在SSL模型的ASR微调阶段重点关注模型压缩，而不是预训练。 密集模型和压缩模型在收敛时通常具有非常不同的分布。 当从密集模型的接近最优的预训练权重进行初始化时，我们需要在有限的训练步骤内调整分布以适应压缩模型，这使得它成为一项更具挑战性的任务。 最后，主干大小（2B 与 10-100M 参数）和数据大小（$\sim$100 万小时对比 $\sim$1,000 小时）比之前的研究要大得多。

3方法

3.1USM骨干架构

在本研究中，我们使用具有 2B 参数的最先进的 USM-CTC [15] 主干网。 编码器由 32 个 Conformer 层 [31] 组成，维度为 1536。 我们在自注意力层中使用相对注意力[32]，有16个注意力头。 深度卷积的内核大小设置为 5。 该模型使用 BEST-RQ [13] 对从 YouTube 收集的 500 多种语言的超过 1200 万小时的语音数据进行了预训练。 BEST-RQ 以 BERT 训练方式运行，它将音频作为输入并预测掩蔽语音表示。 此外，每层左右注意力上下文设置为 128 帧。

给定一个预先训练的 USM，我们提取编码器并将其与对应于单词片段模型 (WPM) 的随机初始化的 softmax 层配对。 此后，模型根据普通联结主义时间分类 (CTC) [33] 损失进行微调。 虽然 RNN-T [3] 或 LAS [34] 由于额外的语言建模能力可能会导致 WER 的改进，但这些模型在推理过程中以自回归方式运行，在推理过程中很难并行化，并且使用如此庞大的编码器具有更高的延迟。 相比之下，基于 CTC 的模型具有仅编码器架构，没有任何自回归依赖性，可以轻松并行化，因此对于大规模模型来说更有效。

3.2 训练本机量化感知 (QAT)

假设我们有一个单线性层矩阵乘法$𝐘=𝐗\otimes 𝐖$，其中${𝐗}^T\in {ℝ}^I$,${𝐘}^T\in {ℝ}^J$和$𝐖\in {ℝ}^{I\times J}$ 分别是输入、输出和权重。 使用每通道权重量化运行矩阵乘法可以表示为：

$${𝐘}_j={𝐬}_j\cdot \left[𝐗\otimes \text{Quantize}\left({𝐖}_j\right)\right],1\le j\le J,$$ (1)

其中$\text{Quantize}(\cdot )$ 操作的作用是：

$$\text{Quantize}\left({𝐖}_j\right)=\text{round}\left(\frac{{𝐖}_j}{{𝐬}_j}\right),$$ (2)

${𝐬}_j\in ℝ$表示第$j$通道的比例，${𝐖}_j$是$j$第 $𝐖$ 列。通过除以 ${𝐖}_j$ 与整数范围的最大值。 对于 int8 和 int4 量化，我们使用 [24] 中引入的最简单的对称量化（例如，对于 int8 为 [-127, 127]，对于 int4 为 [-7, 7]）。 然而，当精度降低到 2 位时，对称量化未充分利用量化桶（即仅使用三个值）。 因此，按照[25]，我们对int2模型采用非对称量化，以及子通道量化，将通道分为几个组，每个组都有专用的尺度。

在量化感知训练的前向传播过程中，我们应用 eq.(1）到模型的所有全连接层，并使用 STE 绕过反向传播时不导数的舍入函数（几乎到处都是零导数）。 更重要的是，我们从 eq.(2) 到 本地 整数类型。 与常用的“假”量化[35]相比，在训练时使用浮点运算，在推理时使用整数运算，这避免了训练和推理之间运算不匹配导致的数值差异。

3.3 基于幅度的修剪，具有$N$:$M$稀疏性

图。 1： 权重矩阵上具有 $N$:$M$ 稀疏度的基于幅度的剪枝图示。 此示例有 $N=2$ 和 $M=4$。

$N$:$M$稀疏性有一个模式：对于每组$M$连续权重，最多有 $N$ 非零值。 在这项工作中，我们重点关注以下常用案例 $M=4$，但它可以轻松扩展到具有任意$M$的模式。如图1所示，在剪枝步骤中，我们首先将稠密权重矩阵重塑为$𝐕\in {ℝ}^{K\times M}=\text{Reshape}\left(𝐖\right)$，其中$K$ 是组数。 然后，我们识别 $N$-每个组的最大幅度权重${\varphi }_k$，并生成二进制掩码$𝐌\in {\{0,1\}}^{K\times M}$ 作者：

(3)

重塑后的权重$𝐕$ 被掩码修剪为：

$$\text{Prune}\left(𝐕\right)=𝐕\odot 𝐌$$ (4)

哪里$\odot$ 表示逐元素乘积。 最后，我们将修剪后的稀疏权重重塑回原始形状。

我们研究了一次性和几次 训练中的修剪。 一次性剪枝仅在 USM 微调阶段开始时更新一次掩模，然后我们冻结掩模并仅调整权重。 少样本修剪更新了掩模 $T_p$ 开始时多次，同样，之后将面膜冷冻。 我们在少样本剪枝中不启用[26]中提出的STE，并且每次迭代时仅更新未剪枝的权重。

3.4 量化和稀疏性联合优化

根据经验，简单地单独应用量化或稀疏性和高压缩比会引入不可避免的回归（例如，参见5.1 和 5.2 节中 2 位量化和 1:4 稀疏度的结果）。 为了在保留WER的同时最大限度地减小模型大小，我们建议从参数精度和矩阵拓扑方面联合压缩模型，并结合量化和稀疏性。 我们提出的方法采用修剪和量化的方式，这使得它在实现过程中更加方便。 如节中所述 3.3，剪枝权重设置为零，直接映射到对称量化的零点，对计算尺度没有影响。 整体训练流程如Algorithm所示 1.

算法1 建议用于 USM 微调的量化和稀疏性感知训练过程。

1:Inputs: Speech-text pairs ($𝐗$, $𝐘$), USM model $f(𝐗;𝐖)$, CTC loss $L(\cdot )$, total training steps $T$, total pruning steps $T_p$

2:Initialization: $𝐖={𝐖}_0$ from BEST-RQ.

3:for $t$ in $1,2,\mathrm{\dots },T$ do

4: if then

5: Update masks for each weight matrix of $𝐖$ with eq.(3)

6: end if

7: Prune each weight matrix of $𝐖$ through the mask with eq.(4)

8: Quantize each weight matrix of $𝐖$ with eq.(2) and run matrix multiplications with eq.(1)

9: Compute CTC loss: $L(f(𝐗;𝐖),𝐘)$

10: Update $𝐖$ with gradients $\frac{\partial L}{\partial 𝐰}$

11:end for

12:Return: Trained model $f(𝐗;𝐖)$

4实验设置

4.1数据集

我们在大规模语音搜索任务上评估所提出的技术。 预训练数据集在章节中描述 3.1. 在微调过程中，我们使用来自语音搜索的 120 万小时美国英语音频文本对的内部数据集来训练模型。 所有数据都是匿名的，我们的数据处理遵循 谷歌人工智能原理 [36]。 数据集的一小部分是手工转录的，其余部分是使用 6 亿参数的教师模型[37]进行伪转录。 在评估中，我们报告了代表语音搜索流量的 8,884 条匿名和手写话语的 WER。

4.2实现细节

除了3.1小节介绍的USM主干架构之外，网络采用 128 维对数梅尔滤波器组能量的输入序列，从 32ms 窗口和 10ms 偏移中提取。 输入特征通过两个二维卷积层下采样 4 倍，并通过全连接层投影到 1,536 维，然后馈送到编码器。 每个卷积层都有一个 $3\times 3$内核，$2\times 2$ 步长分别为 128、32 个通道。

所提出的方法在 Pax 中实现 与图层库 Praxis. 我们不对卷积层和最终的 softmax 层进行量化和稀疏化，因为它们的参数数量要少得多。 这些模型在张量处理单元 (TPU) v3-128 [38] 上进行训练，并使用 Transformer 学习率计划 [39] 的 Adam 优化。 输入特征和编码器模块的基本学习率乘数为 0.5，softmax 层的基本学习率乘数为 2.0。 这些模块的预热步骤分别设置为 5,000 和 1,500。 所有模型均经过 200,000 个步骤的训练，批量大小为 2,048。 在评估过程中，我们 不要 包括用于重新评分的任何其他语言模型，以更好地衡量压缩技术的有效性。

5实验

表1： 量化消融研究的结果。 模型大小比率计算为估计模型大小相对于B0的比率。 PTQ是指训练后量化。

Exp	Model	Voice Search	Model Size
		WER	Ratio
B0	float32 dense 2B CTC USM	4.1	-
E0	int8 PTQ	4.2	25.0%
E1	int8 QAT	4.2	25.0%
E2	int4 PTQ	86.7	12.5%
E3	int4 QAT	4.3	12.5%
E4	int2 QAT	99.9	6.3%
E5	int2 QAT + 16 sub-channel	45.2	7.3%
E6	int2 QAT + 32 sub-channel	32.0	8.3%
E7	int2 QAT + 64 sub-channel	12.3	10.4%

我们进行了三组实验来评估我们提出的方法：1）量化的消融研究； 2）稀疏性消融研究； 3) 所建议组合的总体性能。 通过量化和稀疏性的消融研究，我们的目的是检查和表征这两种技术在微调 USM 模型方面的有效性。 在评估我们提出的组合时，我们与几个密集/压缩基线进行比较，并展示了从所提出的方法获得的最先进的质量和尺寸减小。 请注意，量化/稀疏模型的大小是 估计的 不考虑硬件所需的额外偏移，因此实际模型尺寸会稍大。

5.1 量化消融研究

表1 显示了 int8、int4 和 int2 量化的消融研究结果。 实验 B0 具有 float32 密集 2B 参数 CTC USM 的性能，它作为两项消融研究的上限。 除了 QAT 之外，我们还探讨了 int8 和 int4 的训练后量化（PTQ）的性能。 从结果可以看出，int8 PTQ (E0)和 QAT (E1) 都可以保留 float32 模型的 WER。 对于 int4 精度，PTQ (E2)导致更显着的回归，而 QAT(E3）仅对 float32 模型有边际 WER 回归（4.3 与 4.1）。 当涉及到 int2 时，极端的量化设置、普通的非对称量化不会产生合理的 WER。 随着子通道数量的增加，WER可以逐渐提高到64个子通道时的12.3（E7)。然而，子通道尺度引入了额外的参数（即 float32 每个子通道的参数），分别相当于 16、32 和 64 个子通道模型上原始模型大小的 1%、2% 和 4.1%。 这些观察结果大多与[24, 25]的结论相对应。

5.2稀疏性消融研究

同样，我们在表中显示了稀疏性的消融研究结果2. 具体来说，我们研究 2:4 和 1:4 结构化稀疏模式以及一次性和少数（1k）次修剪计划。 稀疏模型需要额外的 1 位（即二进制）参数来存储掩码，这相当于原始模型大小的 3.1%。 从表中可以看出，一次性（E8)和 1k 次拍摄(E9）与稠密模型相比，2:4 稀疏度的剪枝具有边际回归，其大小为原始模型大小的 53.1%。 当以 1:4 的稀疏度更积极地剪枝模型时，两个模型都无法保留密集的 WER，同时 1k-shot (E11)剪枝实现了比一次性(E10）。 一种可能的解释是，1:4 稀疏度对掩模更敏感，因为它比 2:4 稀疏度是一项艰巨的任务，因此，在这种情况下，少样本更新可以识别更有效的掩模。

表2： $N$ 的消融研究结果：$M$ 稀疏性。 模型大小比率计算为估计模型大小相对于B0的比率。

Exp	Model	Voice Search	Model Size
		WER	Ratio
B0	float32 dense 2B CTC USM	4.1	-
E8	2:4 sparsity one-shot	4.4	53.1%
E9	2:4 sparsity 1k-shot	4.3	53.1%
E10	1:4 sparsity one-shot	11.7	28.1%
E11	1:4 sparsity 1k-shot	10.6	28.1%

5.3量化与稀疏性相结合的整体性能。

根据第5.1和5.2节的消融研究，我们观察到单独的量化或稀疏性在极端压缩比下都无法获得合理的 WER。 这促使我们研究通过我们提出的训练方案结合这两种技术。 此外，我们使用 1B (B1)、600M(B2)和 300M(B3）参数供参考。 桌子 3 显示了该实验的结果。 我们提出的两个具有 int4 量化和 2:4 稀疏度的系统达到了 4.4（E12)和 4.5(E13) 分别为WER。 与基线相比 B0，模型尺寸缩小了10倍以上(B0的9.4%）），相对 WER 回归率仅为 7.3%（4.4 比 4.1）。 与 float32 密集基线相比（B1 到 B3)，E12 的 WER 甚至比1B参数的模型还要好，但它的大小只有1B模型的五分之一左右。 最后，与量化相比（E7)或稀疏(E11)模型在极端压缩率下，我们提出的方法显着降低了两个模型的 WER更接近密集模型B0，同时享受更小的模型尺寸（9.4% vs. 10.4%）。 关于一次性剪枝和少样本剪枝之间的比较，我们没有发现稀疏度为 2:4 的 WER 存在显着差异，这与我们在 5.2. 总之，这些结果证实了我们的主张，即从参数精度和矩阵拓扑方面联合压缩模型比单独的技术更有效。

表3： 所提出的结合量化和 $N$ 范例的结果：$M$ 稀疏性。 此处还提供了不同模型尺寸的基线 USM 结果以供比较。 模型大小比率计算为估计模型大小相对于B0的比率。

Exp	Model	Voice Search	Model Size
		WER	Ratio
B0	float32 dense 2B CTC USM	4.1	-
B1	float32 dense 1B CTC USM	4.5	50.2%
B2	float32 dense 600M CTC USM	4.7	33.5%
B3	float32 dense 300M CTC USM	5.0	18.9%
E7	int2 QAT + 64 sub-channel	12.3	10.4%
E11	1:4 sparsity 1k-shot	10.6	28.1%
E12	int4 QAT + 2:4 sparsity one-shot	4.4	9.4%
E13	int4 QAT + 2:4 sparsity 1k-shot	4.5	9.4%

5.4 限制和讨论

尽管我们从上述实验中获得了非常令人鼓舞的结果，但我们观察到这项工作的一些局限性，我们想在这里讨论并在未来的工作中进一步探讨。 首先，我们注意到这项工作中对 int4 PTQ 和 int2 QAT 有更多的回归。 一个潜在的原因可能是我们在这项工作的微调过程中没有启用全局变分噪声（VN）[40]，因为尚未系统地研究 USM 微调。 相比之下，[25] 已启用它，这已被证明可以提高低位量化性能[41]。 其次，我们在模型剪枝期间没有启用 STE，这可能会提高模型的性能 $N:M$ 稀疏性。 在未来的工作中，我们将研究这两种方法以及已被证明在量化方面有效的其他技术（例如，可学习的尺度、离群值裁剪[42]等）和 $N$:$M$ 稀疏性（例如 SR-STE [26])，以寻求更好的量化和稀疏模型。 借助先进的量化和稀疏模型，我们将研究更激进的组合，例如 int2 加 2:4，以将模型大小推向边缘。

6 结论

USM 显着提高了 ASR 以及其他下游任务的质量并简化了生产。 在本文中，我们首次提出了一种基于 USM 的 ASR 微调的联合量化和稀疏感知范式。 通过对大规模数据集进行广泛的实验和消融研究，我们首先对低位量化和 $N$:$M$ USM 微调上的结构化稀疏性，验证了 int4 QAT 和 2:4 稀疏性的有效性。 更重要的是，结果表明我们提出的量化和稀疏性组合可以进一步将模型大小减少到原始模型的 9.4%，但代价是边际性能回归。 我们还对每种压缩技术的结果进行了系统分析，并讨论了未来研究的局限性和潜在解决方案。

7致谢

我们要感谢 Jeremiah Willcock、Emmanuel Guzman、Xingyu Cai、Wonpyo Park、Peiran Li 和 Suvinay Subramanian 在这项工作中的讨论和帮助。