推理速度22.3倍提升,北航、字節(jié)跳動提出二值化關鍵字識別模型
近年來,以 FSMN 為代表的語音關鍵字識別(KWS)模型在各類邊緣場景得到廣泛應用。然而,語音喚醒應用的實時響應需求和邊緣設備上有限計算與能耗資源間的矛盾一直存在,這阻礙了 KWS 模型在真實世界硬件設備上的部署。近日,人工智能頂會 IJCAI 2022 接收論文結果已經正式公布,北航劉祥龍教授團隊和字節(jié)跳動 AI Lab 智能語音團隊聯(lián)合提出了首個針對 KWS 任務的二值神經網(wǎng)絡 BiFSMN,并在 ARM 設備上實現(xiàn)了高達 22.3 倍和 15.5 倍的推理加速和存儲節(jié)省。
論文地址:https://arxiv.org/pdf/2202.06483.pdf1.
引言
1.1 背景
目前深度神經網(wǎng)絡的卓越性能,主要依賴于高端圖形處理單元進行計算訓練。訓練好的模型,其體積和參數(shù)量通常較大,因此需要大量的存儲空間占用,并且有足夠的計算單元來提高模型運行效率。這導致模型較難部署在手機等部分算力有限、存儲空間有限的邊緣設備上,這也限制了神經網(wǎng)絡的適用場景和部署平臺。
1.2 問題
盡管目前模型二值化取得了進展,但通過現(xiàn)有方法對 KWS 網(wǎng)絡進行二值化仍然遠非理想。大多數(shù)量化方法在推理過程中使用浮點算法,且直接量化會帶來嚴重的性能下降。首先,由于使用了 1 位參數(shù),二值化網(wǎng)絡的表示空間極其有限,難以優(yōu)化。其次,KWS 的現(xiàn)有架構具有固定的模型規(guī)模和拓撲結構,無法在運行時自適應地平衡資源預算。此外,現(xiàn)有的部署框架在現(xiàn)實世界的硬件上實現(xiàn)時還遠未達到二值化網(wǎng)絡的理論加速上限。
1.3 成果
- 提出了一種用于 KWS 的準確且極其高效的二元神經網(wǎng)絡 BiFSMN。
- 構建了 HED,方案來強調高頻信息,以優(yōu)化二值化網(wǎng)絡的訓練。
- 提出了 TBA,以在運行時實現(xiàn)即時和自適應的精度 - 效率權衡。
BiFSMN 通過令人信服的精度改進優(yōu)于現(xiàn)有的二值化方法,甚至可以與全精度對應物相媲美。此外,該研究在 ARMv8 實際設備上的 BiFSMN 實現(xiàn)了 22.3 倍加速和 15.5 倍存儲節(jié)省。
2. 方法概述
2.1 HED —— 高頻增強蒸餾
該研究發(fā)現(xiàn),信息傾向于邊緣的本質是:基本的二值化表示傾向于集中在高頻分量上。該研究使用 2D Haar 小波變換 (WT) [Meyer, 1992],其常用作分離水平、垂直邊緣的可分離變換,將表示分解為低頻和高頻分量。輸入到特定層的隱藏狀態(tài) H 可以表示為小波函數(shù)族的加權和,
其中 是具有特定時間參數(shù)的母小波函數(shù),
是分辨率級別, 和 K 確定波形的平移,為了測量表示的單個分量所傳達的信息量,使用相對小波能量來定義信息量[Rosso et al., 2001]。第
層的小波能量
首先計算為:
當通過一次分解得到低頻和高頻系數(shù)時,它們的相對小波能量
可以表示為:
較大的相對小波能量表明信息更多地聚集在該分量中。如圖所示,與全精度表示相比,二值化表示的高頻分量的相對小波能量顯著增加,這意味著二值化表示向高頻分量傾斜。
基于上述分析,該研究提出了一種用于二值化感知訓練的高頻增強蒸餾方法。該方案利用預訓練的全精度 D-FSMN 作為教師,并在蒸餾過程中增強其隱藏層特征的高頻分量。具體來說,該研究對原始特征應用小波變換,去除低頻分量,然后對高頻分量應用小波逆變換(IWT)函數(shù)來恢復選定的特征。該過程可以表述如下:
然后將強調的高頻表示添加到原始表示中:
其中 是標準偏差。受 [Martinez et al., 2020] 的啟發(fā),該研究最小化了來自教師的
和直接來自學生的隱藏層特征
之間的注意力蒸餾損失,表示為:
其中 表示第
個塊,
是 L2 范數(shù)。
上面的 HED 方案使二值化學生網(wǎng)絡更容易利用強調的全精度表示中的基本信息并提高準確性。2.2 TBA —— 可細化二值結構該研究提出了一種用于 KWS 的 Thinnable Binarization Architecture (TBA),它可以在運行時選擇具有更少層的更薄模型,從而直接減少計算消耗。該研究把包含 N 個塊的基本二值化架構的整個主干網(wǎng)絡
表示為:
其中 分別是二值化網(wǎng)絡和
二值化 D-FSMN 塊,
是網(wǎng)絡的輸入。
形成的 TBA 結構可以定義為:
其中 是所選層的間隔,取值僅限于可整除 N。每個可細化塊
可以定義為:
可細化網(wǎng)絡架構將通過用恒等函數(shù)替換中間塊來跳過每個 層的中間塊, 下圖顯示了該研究的可細化二值化架構的形式化。
此外,該研究還提供了 的實例,如圖所示。
為了優(yōu)化提議的 TBA 的二值化感知訓練,該研究采用統(tǒng)一層映射策略來更好地對齊和學習 HED 中的表示:
來自不同分支的梯度在反向傳播過程中累積以共同更新權重。根據(jù)可細化架構中的壓縮比,加權損失可以計算為:
其中 表示
的交叉熵損失,
是控制蒸餾影響的超參數(shù)。
2.3 FBCK —— 用于高效硬件部署的快速按位計算內核
為了在計算資源有限的邊緣設備上高效部署,該研究通過新的指令和寄存器分配策略進一步優(yōu)化 1 位計算,以加速邊緣設備上廣泛使用的 ARMv8-A 架構的推理。該研究稱之為快速按位計算內核 (FBCK)。
根據(jù) ARMv8 架構上的寄存器數(shù)量,該研究首先將內核中的寄存器重新分配為五個分區(qū),以提高寄存器利用率并減少內存占用:分區(qū) A 有四個寄存器(寄存器 v0 除外)用于一個輸入(權重 / 激活),B 有兩個用于另一個輸入,C 有 8 個用于 EOR 和 CNT 的中間結果,D 有 8 個用于一個循環(huán)中的輸出,E 有 8 個用于最終結果。每個輸入都打包為 INT16。A 中的每個寄存器存儲一個輸入并重復 8 次,而 B 中的每個寄存器存儲 8 個不同的輸入。該研究先對 A 和 B 的一個寄存器進行 EOR 和 CNT,得到 32 個 INT8 結果到中間分區(qū) C,然后執(zhí)行 ADD 將 INT8 累加到 D,對 B 的另一個寄存器做同樣的事情。經過 16 次循環(huán),最后,該研究使用長指令 ADALP 將存儲在 D 中的 INT8 數(shù)據(jù)累積到一個 INT16 寄存器(在 E 中),它將 INT8 數(shù)據(jù)擴展為雙倍寬度。FBCK 在計算過程中充分利用了寄存器,幾乎沒有空閑位。
3. 實驗
在本章,該研究從推斷精度、理論計算、部署性能等角度對 BiFSMN 進行全面評估。實驗證明,BiFSMN 在各個方面大幅領先現(xiàn)有的二值神經網(wǎng)絡,并展現(xiàn)了在實際部署場景下的出眾潛力。
3.1 對比驗證
該研究首先進行消融研究,以調查所提出的高頻增強蒸餾 (HED) 和可細化二值化架構 (TBA)在 D-FSMN 上對 Speech Commands V1-12 和 Speech Commands V2-12 KWS 任務的影響,包括高頻增強蒸餾 (HED) 和可細化二值化架構 (TBA)。
表 1 中結果表明,普通二值化基線方法在兩個數(shù)據(jù)集中都出現(xiàn)了顯著的性能下降。樸素的蒸餾方案 Distill 有助于提高基本 D-FSMN 架構的準確性,而 HED 的應用提高了基于蒸餾的性能。聯(lián)合使用 HED 和 TBA 進一步縮小了二值化模型和全精度模型之間的準確率差距,最終使得在這兩個數(shù)據(jù)集上的準確率均小于 3%。
其次將 BiFSMN 與現(xiàn)有的結構無關二值化方法進行比較,包括 BNN [Courbariaux et al., 2016]、DoReFa [Zhou et al., 2016]、XNOR [Rastegari et al., 2016]、Bi-Real [Liu et al., 2018]、IR-Net [Qin et al., 2020] 和 RAD [Ding et al., 2019]。
表 2 中結果表明,該研究的 1 位 BiFSMN 完全優(yōu)于其他 SOTA 二值化方法。值得注意的是,BiFSMN在兩個數(shù)據(jù)集上的平均準確率下降僅為4%,并遠超其他二值化方法。
其次,為了從架構的角度驗證 TBA 的優(yōu)勢,該研究還將其與 KWS 中廣泛使用的各種網(wǎng)絡進行了比較,包括 FSMN [Zhang et al., 2015]、VGG190 [Simonyan and Zisserman, 2014]、BCResNet [Kim et al.,2021] 和 Audiomer [Sahu al.,2021]。該研究使用 XNOR 和 IR-Net 對這些架構進行二值化。
表 3 表明, HED 通??蓱糜诨陬愃?FSMN 的架構,并對二值化模型性能產生影響。此外,配備 TBA 的 BiFSMN 可以進一步在運行時在準確性和效率之間取得平衡。例如,更薄的版本 BiFSMN0.5× with 4 blocks 和 BiFSMN0.25× with 2 blocks 甚至在 Speech Commands V1-12 任務上實現(xiàn)了 23.8× 和 28.8× FLOPs 節(jié)省,而不犧牲準確性(僅下降 0.16% 和 0.13%)。
該研究進一步修剪模型寬度并提供一個極小的 BiFSMNS(具有 32 個主干內存大小和 64 個隱藏大?。挥?0.05M 參數(shù)和 9.16M FLOP,證明該研究的方法在微型網(wǎng)絡上也能很好地工作。
3.2 部署效率
除了推理精度與理論計算性能,在現(xiàn)實世界的邊緣設備上運行時,KWS 任務對于低內存占用和快速實時響應具有迫切需求。為了驗證 BiFSMN 的實際部署效率,該研究在 1.2GHz 64 位 ARMv8 CPU Cortex-A53 的 Raspberry Pi 3B + 上測試了 BiFSMN 的實際速度。
如圖 5 所示,由于提出了優(yōu)化的 1 位快速按位計算內核, BiFSMN 與全精度對應物相比提供了 10.9 倍的加速度。它也比現(xiàn)有的開源高性能二值化框架(如 daBNN 和 Bolt)快得多。此外,受益于可細化架構,BiFSMN 可以根據(jù)設備上的資源自適應地平衡運行時的準確性和效率,切換到分別帶來了 15.5× 和 22.3× 加速。這表明 BiFSMN 可以在實際推理中滿足不同的資源約束。