该研究提出了首个针对 KWS 任务的二值神经网络 BiFSMN,通过令人信服的精度改进优于现有的二值化方法,甚至可以与全精度对应物相媲美。此外,该研究在 ARMv8 实际设备上的 BiFSMN 实现了 22.3 倍加速和 15.5 倍存储节省。
近年来,以 FSMN 为代表的语音关键字识别(KWS)模型在各类边缘场景得到广泛应用。然而,语音唤醒应用的实时响应需求和边缘设备上有限计算与能耗资源间的矛盾一直存在,这阻碍了 KWS 模型在真实世界硬件设备上的部署。近日,人工智能顶会 IJCAI 2022 接收论文结果已经正式公布,北航刘祥龙教授团队和字节跳动 AI Lab 智能语音团队联合提出了首个针对 KWS 任务的二值神经网络 BiFSMN,并在 ARM 设备上实现了高达 22.3 倍和 15.5 倍的推理加速和存储节省。
论文地址:https://arxiv.org/pdf/2202.06483.pdf
1. 引言
1.1 背景
目前深度神经网络的卓越性能,主要依赖于高端图形处理单元进行计算训练。训练好的模型,其体积和参数量通常较大,因此需要大量的存储空间占用,并且有足够的计算单元来提高模型运行效率。这导致模型较难部署在手机等部分算力有限、存储空间有限的边缘设备上,这也限制了神经网络的适用场景和部署平台。
1.2 问题
尽管目前模型二值化取得了进展,但通过现有方法对 KWS 网络进行二值化仍然远非理想。大多数量化方法在推理过程中使用浮点算法,且直接量化会带来严重的性能下降。
首先,由于使用了 1 位参数,二值化网络的表示空间极其有限,难以优化。其次,KWS 的现有架构具有固定的模型规模和拓扑结构,无法在运行时自适应地平衡资源预算。此外,现有的部署框架在现实世界的硬件上实现时还远未达到二值化网络的理论加速上限。
1.3 成果
- 提出了一种用于 KWS 的准确且极其高效的二元神经网络 BiFSMN。
- 构建了 HED,方案来强调高频信息,以优化二值化网络的训练。
- 提出了 TBA,以在运行时实现即时和自适应的精度 - 效率权衡。
BiFSMN 通过令人信服的精度改进优于现有的二值化方法,甚至可以与全精度对应物相媲美。此外,该研究在 ARMv8 实际设备上的 BiFSMN 实现了 22.3 倍加速和 15.5 倍存储节省。
2. 方法概述
2.1 HED —— 高频增强蒸馏
该研究发现,信息倾向于边缘的本质是:基本的二值化表示倾向于集中在高频分量上。该研究使用 2D Haar 小波变换 (WT) [Meyer, 1992],其常用作分离水平、垂直边缘的可分离变换,将表示分解为低频和高频分量。
输入到特定层的隐藏状态 H 可以表示为小波函数族的加权和,
其中 是具有特定时间参数的母小波函数, 是分辨率级别, 和 K 确定波形的平移,为了测量表示的单个分量所传达的信息量,使用相对小波能量来定义信息量[Rosso et al., 2001]。第层的小波能量 首先计算为:
当通过一次分解得到低频和高频系数时,它们的相对小波能量可以表示为:
较大的相对小波能量表明信息更多地聚集在该分量中。如图所示,与全精度表示相比,二值化表示的高频分量的相对小波能量显着增加,这意味着二值化表示向高频分量倾斜。
基于上述分析,该研究提出了一种用于二值化感知训练的高频增强蒸馏方法。该方案利用预训练的全精度 D-FSMN 作为教师,并在蒸馏过程中增强其隐藏层特征的高频分量。具体来说,该研究对原始特征应用小波变换,去除低频分量,然后对高频分量应用小波逆变换(IWT)函数来恢复选定的特征。该过程可以表述如下:
然后将强调的高频表示添加到原始表示中:
其中 是标准偏差。受 [Martinez et al., 2020] 的启发,该研究最小化了来自教师的 和直接来自学生的隐藏层特征 之间的注意力蒸馏损失,表示为:
其中 表示第 个块,是 L2 范数。
上面的 HED 方案使二值化学生网络更容易利用强调的全精度表示中的基本信息并提高准确性。
2.2 TBA —— 可细化二值结构
该研究提出了一种用于 KWS 的 Thinnable Binarization Architecture (TBA),它可以在运行时选择具有更少层的更薄模型,从而直接减少计算消耗。该研究把包含 N 个块的基本二值化架构的整个主干网络表示为:
其中 分别是二值化网络和二值化 D-FSMN 块,是网络的输入。形成的 TBA 结构可以定义为:
其中 是所选层的间隔,取值仅限于可整除 N。每个可细化块 可以定义为:
可细化网络架构将通过用恒等函数替换中间块来跳过每个 层的中间块, 下图显示了该研究的可细化二值化架构的形式化。
此外,该研究还提供了 的实例,如图所示。
为了优化提议的 TBA 的二值化感知训练,该研究采用统一层映射策略来更好地对齐和学习 HED 中的表示:
来自不同分支的梯度在反向传播过程中累积以共同更新权重。根据可细化架构中的压缩比,加权损失可以计算为:
其中 表示 的交叉熵损失,是控制蒸馏影响的超参数。