原創 Synced 機器之心
機器之心發布
作者:王瀚銳等
運用注意力機制的 Transformer 模型近幾年在 NLP 領域獲得了廣泛應用。然而,由於參數量和計算量巨大,Transformer 模型難以在存儲和算力有限的邊緣硬體設備上高效部署。為了解決 Transformer 的低效問題,來自 MIT 的研究人員提出了 HAT: Hardware-Aware Transformers,針對不同的硬體設備的特性,為每個硬體搜索出一個高效的 Transformer 模型,從而在保持精確度的前提下大幅降低內存消耗。在同樣的精度下,相比於基線 Transformer, HAT 可以獲得 3 倍加速,3.7 倍模型壓縮。
該論文已被自然語言處理頂會 ACL 2020 收錄。此外,HAT 的所有代碼和模型已經在 GitHub 上開源,作者也將在 7 月 8 日 / 9 日的 ACL 大會上線上宣講他們的工作。
論文連結:https://arxiv.org/abs/2005.14187
GitHub:https://github.com/mit-han-lab/hardware-aware-transformers
B站介紹: https://www.bilibili.com/video/BV1mt4y197FL/
近年來,自然語言處理領域受 Transformer 的驅動獲得了快速發展,Transformer 模型也被廣泛應用於多種任務,例如機器翻譯、聊天機器人、文本摘要等等。然而,Transformer 的高準確性需要非常高的模型參數量和計算量來支撐,這對於受到存儲大小、算力和電池容量限制的邊緣計算設備(比如手機和物聯網硬體)來說是很大的挑戰。例如,如果在樹莓派上使用 Transformer-Big 模型,它需要運行 20 秒才可以完成一個 30 詞長度的句子的翻譯,這對於很多需要實時對話和反饋的場景來說是無法接受的。
為了獲得更高效和快速的 Transformer 模型,MIT 等機構的研究者提出了 HAT: Hardware-Aware Transformers,藉助神經網絡搜索(NAS)技術,在搜索過程中加入硬體反饋,來對每一個硬體平臺設計一個專用的高效 Transformer 網絡結構。
在搜索空間上,他們提出打破之前設計 Transformer 的兩點陳規。首先,傳統 Transformer 的所有解碼層都從最後一層編碼層獲得輸入,這樣形成了一個從編碼器到解碼器的信息瓶頸,對小模型尤其不友好;對此,研究者提出「任意編碼器 - 解碼器注意力」(Arbitrary Encoder-Decoder Attention),允許解碼器獲得多個和任意某些個編碼器的輸出。從而使得編碼器不同的抽象層的信息都可以被解碼器獲取和使用。另外,傳統 Transformer 的所有層都有相同的網絡結構,研究者提出「異構層」(Heterogenous Layers)來使得每層都可以有不同的隱藏層維度(Hidden Dim)和注意力頭數 (Head Number),編 / 解碼器也可以有不同的詞向量長度 (Embedding Dim) 和層數等等。
在搜索算法上,為了能夠減少搜索開銷,實現環保 AI,他們採用權重共享的方法來訓練出一個母網絡 SuperTransformer 來涵蓋在搜索空間中的全部可能模型,這樣一來,其中的每個子網絡 SubTransformer 可以直接繼承母網絡中對應部分的權重,快速估計子網絡的精確度。之後,研究者使用進化搜索(Evolutionary Search),利用預先設定好的在目標硬體上的運行時間來作為限制,在母網絡中搜索出一個滿足運行時間並且精度夠高的子網絡,這就是他們想要得到的為目標硬體設計的專用 SubTransformer 網絡。
研究者在四個機器翻譯任務(WMT『14 EN-De(英語 - 德語)、WMT』14 EN-Fr(英語 - 法語)、WMT』19 EN-DE(英語 - 德語)和 IWSLT『14 De-En(德語 - 英語)),以及三個不同的硬體平臺(樹莓派、Intel CPU 和 Nvidia GPU)上驗證了論文所提方法的有效性。
在同樣的精度下,相比於基線 Transformer, HAT 可以獲得 3 倍加速,3.7 倍模型壓縮;相比於 Evolved Transformer,HAT 有 2.7 倍的加速和 3.6 倍的模型壓縮,並且將搜索開銷降到了 1.2 萬分之一。同時,HAT 的方法也與其他模型壓縮技術兼容,例如,研究者將搜索得到的 HAT 模型進行了 4-bit 壓縮,進一步獲得了僅為基線 Transformer 1/25 大小的模型。
圖 1:HAT 為每個硬體設計一個高效專用網路。
研究背景
在 Transformer 的性能測試中,研究者發現了兩個常見且重要的陷阱:
第一,計算量(FLOPs)不能反映真正的運行時間。例如,在下圖中,圈出的兩個模型具有相同的計算量,但是不同的隱藏層大小和層數,導致了兩個模型的運行時間有 5 倍的差距。
圖 2:計算量(FLOPs)不能反映真正的運行時間。
第二,不同的硬體需要有不同的高效模型設計策略。例如,在下圖中,隱藏層的大小顯著地影響樹莓派的運行時間,但是對 GPU 的運行時間幾乎沒有影響。
基於以上兩點,研究者提出將硬體運行時間反饋到模型設計中來,並且為每個硬體平臺設計一個專用 Transformer 架構。
圖 3:不同的硬體需要有不同的高效模型設計策略。
Hardware-Aware Transformers
搜索空間
通過打破傳統 Transformer 設計的兩個陳規,作者們構建了一個很大的搜索空間。
1. 任意編碼器 - 解碼器注意力(Arbitrary Encoder-Decoder Attention)
不同的編碼器層抽取不同抽象層的輸入信息,在基線 Transformer 中,所有的解碼層都必須以編碼器的最後一層作為輸入。這使得編碼器和解碼器之間形成了一個信息瓶頸,並強制所有的解碼層只從高度抽象層解碼信息而忽略了低層信息,對小模型設計不友好。
研究者提出 「任意編碼器 - 解碼器注意力」 來打破信息瓶頸,並且學出編碼器和解碼器之間最合適的連接關係。每一個解碼器層可以注意(attend)到任意且多個編碼器層。從多個編碼器層輸出的 Key 和 Value 向量在句子長度的維度上進行拼接,並且被送到解碼器中的「交叉注意力」(Cross-Attention)模塊。因為這個機制沒有引入多餘的參數量,所以沒有多餘的內存開銷。另外,它對運行時間的影響也幾乎可以忽略。例如,當每個解碼器層都注意兩個編碼器層時,Transformer-Big 模型在 GPU 上的運行時間僅僅升高了 0.4%。
圖 4: 任意編碼器 - 解碼器注意力。
2. 異構層(Heterogeneous Layers)
傳統的 Transformer 對所有的層重複同樣的網絡結構。在 HAT 中,所有的層都可以有不同的結構,例如不同的頭數(Head Number),隱藏層維度(Hidden Dim),詞向量長度(Embedding Dim)等等。在注意力機制層中,Voita et al. (2019) 提出很多頭是多餘的,因此,研究者讓頭的數量成為彈性變化的(Elastic),每一層可以自行決定必要的頭數。在 FFN 層中,輸入特徵向量被投射到一個更大的維度上(隱藏層維度)並被非線性函數激活。傳統上,隱藏層維度的大小是詞向量長度的 2 倍或者 4 倍。但是,因為每層需要的參數量會因抽取特徵難度的不同而不同,所以固定的 2 倍或 4 倍是不合理的。因此在 HAT 中隱藏層維度也是彈性的。HAT 也支持彈性詞向量長度,但要注意的是,編碼器和解碼器的詞向量長度可以不同,但是在編 / 解碼器內部的層之間,詞向量長度保持一致。另外,編碼器和解碼器的層數也是彈性的。
圖 5: 異構層。
SuperTransformer 母網絡訓練
為了得到高效模型,構建一個足夠大的搜索空間至關重要。然而,完整訓練搜索空間中的子網絡來比較精確度的開銷過大。之前的 SOTA 模型 The Evolved Transformer (So et al. 2019) 就是利用這種方法進行搜索。去年 ACL2019 中,Strubell et al. (2019) 指出了 Evolved Transformer 搜索排放的二氧化碳接近五輛汽車在整個使用壽命中的二氧化碳排放總量,造成了巨大的環境負擔。MIT Technology Review 也對這一點進行了報導。
因此,為了可以實現環保 AI(Green AI),研究者採用 SuperTransformer 母網絡來提供精確度的近似。母網絡可以快速測試搜索空間中的任意子網絡精度,而無需子網絡訓練。SuperTransformer 母網絡是搜索空間中最大的模型,並且通過權重共享(Pham et al. 2018)包含了搜索空間中全部網絡。所有的子網絡 SubTransformer 共享他們在母網絡中的共同部分。例如,所有子網絡 SubTransformer 共享彈性詞向量和隱藏層權重的前部,共享注意力層中的 Query、Key 和 Value,也共享編 / 解碼器的共同前幾層。
在訓練過程中,研究者均勻地對在搜索空間中的所有子網絡進行採樣,得到梯度並且更新子網絡對應部分的權重。所有的子網絡獲得了等量的訓練並且均可以獨立地完成任務。在實際中,母網絡 SuperTransformer 所需的訓練時間與單個基線 Transformer 相近,所以開銷很低。
圖 6: 母網絡 SuperTransformer 權重共享。
完成訓練後的母網絡是一個快速且精確的模型精度近似器(Performance Proxy)。給定一個子網絡 SubTransformer,可以從母網絡繼承這部分權重,並且快速地測試得到 validation loss。這樣就無需完整訓練子網絡。通過下圖可以看到,近似器是非常精準的。子網絡的 validation loss 越低,最終的完整訓練得到的 BLEU 越高。
圖 7: 母網絡 SuperTransformer 提供精確的子網絡 SubTransformer 精度近似。
進化搜索(Evolutionary Search)子網絡
給定一個目標硬體上的運行時間限制,研究者使用進化搜索算法來搜索得到子網絡 SubTransformer。
圖 8: 進化搜索(Evolutionary Search)子網絡。
進化搜尋引擎利用 SuperTransformer 來獲得子網絡的 validation loss,並利用一個運行時間預測器來得到子網絡在目標硬體上的運行時間。在每一步中,他們只將運行時間短於時間限制的子網絡加入到種群(Population)中。具體而言,他們採用種群大小 125,母種群大小 25,重組 (Crossover) 種群大小 50,突變 (Mutation) 種群大小 50,0.3 突變機率。
有兩種方法可以獲取硬體運行時間。第一,線上測試法,在搜索的過程中即時測試子網絡的運行時間。然而,這種方法有數個缺點。硬體每次運行的時間受很多因素限制,很不穩定,例如散熱的影響。對於每個模型都需要運行數百次取平均來得到較精確的時間。另外,在進化搜索中,每一步都需要測試種群中的很多個子模型的時間,因此線上測試非常耗時,會成為整個搜索過程的瓶頸。
因此,在 HAT 中,研究者使用線下方法,訓練一個預測器來快速且精確地給出子模型的運行時間。具體來說,他們將子網絡的架構編碼成一個向量,並且採集一個 [網絡架構,運行時間] 數據對的數據集。然後訓練一個多層感知器模型(MLP)來回歸子網絡的運行時間。對於每個硬體,他們採集 2000 個數據點,然後訓練一個三層 MLP,每層維度 400,使用 ReLU 為激活函數。
實驗表明,他們的預測器可以得到很精確的結果,在樹莓派上的 RMSE 僅為 0.1 秒:
圖 9: 運行時間預測器。
最終,他們將搜得的子網絡從頭進行一次完整的訓練,並測試得到最終的模型精度。
實驗
研究者在四個機器翻譯任務和三種硬體平臺上進行了實驗和分析。四個任務為:WMT『14 EN-De(英語 - 德語)、WMT』14 EN-Fr(英語 - 法語)、WMT』19 EN-DE(英語 - 德語)、IWSLT『14 De-En(德語 - 英語);三種硬體為:配備 ARM Cortex-A72 CPU 的樹莓派、Intel Xeon E5-2640 CPU 和 Nvidia TITAN Xp GPU。
與基線 Transformer 對比
在多種平臺和任務中,HAT 相比基線 Transformer 均有更好的精度 - 速度 trade-off 曲線,在相同精度下可取得 3 倍加速和 3.7 倍的模型壓縮。
圖 10:HAT 相比基線 Transformer 有更好的精度 - 速度 trade-off。
專用模型的必要性
兩個搜索得到的模型表現如下。可以看到,GPU 高效模型在 ARM CPU 上並不是最高效的,ARM CPU 高效模型在 GPU 上也不高效。這進一步證明了為不同硬體設計專用模型的必要性。
圖 11:有必要為不同硬體設計專用模型。
與 Evolved Transformer 對比
在樹莓派上運行 WMT『14 EN-Fr(英語 - 法語)任務時,相比 Evolved Transformer,HAT 可以取得 2.7 倍加速,3.7 倍模型壓縮,3.2 倍計算量降低,並節省超過 1 萬倍的搜索開銷。
圖 12:HAT 與 Evolved Transformer 對比。
與 Levenshtein Transformer 和 Lite Transformer 對比
下圖為不同模型在樹莓派上運行 WMT『14 EN-De(英語 - 德語)任務的對比。相比其他模型,HAT 可以獲得最高的 BLEU 和最低的運行時間。值得注意的是,HAT 方法與 Levenshtein 和 Lite Transformer 中提出的新操作具有正交性(Orthogonal),可以結合使用。
圖 13:HAT 與 Levenshtein Transformer 和 Lite Transformer 等對比。
與其他壓縮方法相兼容
下圖為在 WMT』14 EN-Fr(英語 - 法語)任務上,對 HAT 搜索模型進行量化的結果,在 4-bit 量化的情況下,HAT 模型與基線 Transformer 有相近的精度,但是模型大小可壓縮 25 倍。
圖 14:HAT 與通用模型壓縮方法兼容,獲得 25 倍壓縮。
最後,研究者表示,他們將在 ACL 2020 大會上介紹自己的工作,在線問答將在 7 月 8 日北京時間晚 21 點 @13B Machine Translation-15 頻道以及 7 月 9 日北京時間早 5 點 @15B Machine Translation-18 頻道進行,感興趣的老師、同學和工業界同事們可以前去交流。
原標題:《ACL 2020 | 模型壓縮25倍,MIT 韓松組提出高效適配不同硬體的HAT模型》
閱讀原文