模型壓縮25倍,MIT 韓松組提出高效適配不同硬體的HAT模型

2020-12-25 澎湃新聞

原創 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模型》

閱讀原文

相關焦點

  • 模型壓縮95%,MIT韓松等人提出新型Lite Transformer
    該研究是由 MIT 電氣工程和計算機科學系助理教授韓松領導的。韓松的研究廣泛涉足深度學習和計算機體系結構,他提出的 Deep Compression 模型壓縮技術曾獲得 ICLR2016 最佳論文,論文 ESE 稀疏神經網絡推理引擎 2017 年曾獲得晶片領域頂級會議——FPGA 最佳論文獎,引領了世界深度學習加速研究,對業界影響深遠。
  • 學者故事:從深鑑科技創始人到MIT教授,韓松帶你探尋深度壓縮模型世界的奧妙
    他的研究也廣泛涉足深度學習和計算機體系結構,他提出的 Deep Compression 模型壓縮技術曾獲得 ICLR'16 最佳論文,ESE 稀疏神經網絡推理引擎獲得 FPGA'17 最佳論文,對AI算法在移動端的高效部署影響深遠。他的研究成果在 Xilinx、NVIDIA、Facebook、Samsung 得到廣泛應用。
  • 韓松、朱俊彥等人提出GAN壓縮法:算力消耗不到1/9,現已開源
    而且,就目前來看,邊緣設備又多數受到內容以及電池之類硬體層面上的限制,也阻礙了 GAN 在邊緣設備上的部署。 因此,基於 GAN 和 cGAN 在圖像合成領域存在的這些問題,韓松團隊提出了 GAN 壓縮,這是一種通用壓縮方法,用於減少 GAN 的推理時間以及計算成本。
  • 尋找最佳的神經網絡架構,韓松組兩篇論文解讀
    機器之心原創作者:Luo Sainan編輯:Hao Wang、Haojin YangMIT 韓松組一直走在深度模型剪枝,壓縮方向的前沿。最近他們也開始嘗試運用 NAS 神經框架搜索技術來推進研究。為了針對硬體平臺設計專用的高效網絡結構,作者將神經網絡的硬體性能(如延遲)納入到優化目標中。延遲(Latency)是一項不可忽視的重要硬體指標,然而它是不可微的。為了解決這一問題,作者將網絡延遲建模為連續函數並將其優化為正則化損失,如圖 3 所示。此外,作者還提出了一種基於 REINFORCE 的算法作為處理硬體指標的一種替代策略。
  • 剪枝需有的放矢,快手&羅切斯特大學提出基於能耗建模的模型壓縮
    正因為有很多實際需求,快手嘗試構建更高效和有目的性的模型壓縮算法。在本文介紹的這兩篇論文中,他們以能耗為約束構建更高效和合理的壓縮方法。雖然論文以能耗為目標,但實際上可以把它換成推斷時間等不同的度量,因此這類神經網絡壓縮方法有更廣泛的應用。
  • 麻省理工 HAN Lab 提出 ProxylessNAS 自動為目標任務和硬體定製高效 CNN 結構
    在同等的 top-1 準確率下 (74.5% 以上), ProxylessNAS 的手機實測速度是當今業界標準 MobileNetV2 的 1.8 倍。在用 ProxylessNAS 來為不同硬體定製神經網絡結構的同時,我們發現各個平臺上搜索到的神經網絡在結構上有很大不同。這些發現為之後設計高效 CNN 結構提供新的思路。
  • 微軟亞洲研究院的研究員們提出了一種模型壓縮的新思路
    然而,這些模型所包含的參數量巨大,計算成本高昂,極大地阻礙了此類模型在生產環境中的應用。為了解決該問題,來自微軟亞洲研究院自然語言計算組的研究員們提出了一種模型壓縮的新思路。 隨著深度學習的流行,很多大型神經網絡模型誕生,並在多個領域中取得當前最優的性能。尤其是在自然語言處理(NLP)領域中,預訓練和調參已經成為其中大多數任務的新範式。
  • 深度學習模型壓縮方法
    ,為了追求更高的性能;另一派為工程派,旨在將算法更穩定、高效的落地在硬體平臺上,效率是其追求的目標。複雜的模型固然具有更好的性能,但是高額的存儲空間、計算資源消耗是使其難以有效的應用在各硬體平臺上的重要原因。最近正好在關注有關深度學習模型壓縮的方法,發現目前已有越來越多關於模型壓縮方法的研究,從理論研究到平臺實現,取得了非常大的進展。
  • 壓縮26倍,神經網絡仍高效,Facebook新研究造福嵌入式設備
    毫無疑問,現在,嵌入式設備也需要高效的神經網絡加持。但是,如何在嵌入式設備上實現高效的神經網絡,可不是一件簡單的事情。性能,功耗,成本,都是不得不考慮的問題。每一種不同的應用場景,都需要在神經網絡的大小和精確度之間進行特定的權衡(trade-off)。
  • 深度學習模型壓縮與加速綜述
    1、簡介深度學習模型壓縮與加速是指利用神經網絡參數和結構的冗餘性精簡模型,在不影響任務完成度的情況下,得到參數量更少、結構更精簡的模型。被壓縮後的模型對計算資源和內存的需求更小,相比原始模型能滿足更廣泛的應用需求。
  • 只需一個損失函數、一個超參數即可壓縮BERT,MSRA提出模型壓縮新方法
    這些模型從其「過參數化」屬性中獲益,它們往往包含數百萬甚至數十億個參數,這就使得此類模型計算成本高昂,且從內存消耗和高延遲的角度來看其計算是低效的。這一缺陷極大地阻礙了此類模型在生產環境中的應用。為了解決該問題,研究者提出很多神經網絡壓縮技術。一般而言,這些技術可分為三類:量化、權重剪枝和知識蒸餾(KD)。其中,KD 能夠壓縮預訓練語言模型,並因此得到了極大關注。
  • 面對千萬級推薦,如何壓縮模型最高效?這是騰訊看點新框架
    機器之心發布機器之心編輯部推薦領域的模型壓縮一直是個極具挑戰性的問題,面對實際部署中數千萬或億級的推薦項,如何壓縮才能獲得最佳性能?由 SIAT-NLP 組與騰訊看點推薦團隊共同完成的論文《A Generic Network Compression Framework for Sequential Recommender Systems》,提出了一個序列推薦模型框架「CpRec」,壓縮率更高、訓練和推斷速度也更快。目前這篇論文已被 SIGIR 2020 接收。
  • 想在手機上用自己的機器學習模型?谷歌爸爸的模型壓縮包你滿意
    該 ML Kit 開發套件中的一項即將可用的核心功能,是由谷歌 research 團隊開發的 Learn2Compress 技術所驅動的一項自動模型壓縮服務。Learn2Compress 技術能夠在 TensorFlow Lite 內定製移動端深度學習模型,定製後的模型可高效的運行在移動端設備上,無需擔心內存不足和運行速度過慢。
  • [Google]BERT壓縮到7MB!最新基於最優子詞和共享投影的極限語言壓縮模型
    我們的方法能夠將BERT_Base模型壓縮60倍以上,只是稍微降低下遊任務指標,就可以得到內存不足7MB的語言模型。實驗結果還表明,與其他先進的壓縮技術相比,該方法具有更高的壓縮效率和精度。我們進行了測試通用語言建模視角和下遊任務的實驗,證明了這兩類任務都具有高壓縮比的競爭性能。隨著神經網絡本身的普及,對神經網絡模型壓縮的研究也隨之興起,因為這些模型對於當時的硬體來說往往是內存密集型的。NLP應用中的模型壓縮工作大致可分為四類:矩陣逼近、參數剪枝/共享、權重量化和知識蒸餾。
  • 適配難?難適配!AI嵌入式設備部署如何搞?秘訣在此!
    當在實際業務場景中部署AI模型時,會面臨一個部署方案設計的問題:不同的使用場景和業務問題中對模型的部署方式要求也有所差異。 AI模型有著多樣化的部署方式,常見的模型部署形態分為兩種:伺服器部署和嵌入式部署。
  • 綜述論文:當前深度神經網絡模型壓縮和加速方法速覽
    另外,使用遷移/壓縮濾波器的方法只適用於全卷積神經網絡。低秩分解和遷移/壓縮濾波器的方法提供了一種端到端的流程,並且它們很容易直接在 CPU/GPU 環境中實現。而參數修剪和共享使用了不同的方法,如向量量化、二進位編碼和係數約束以執行這些任務,通常他們需要花一些處理步驟才能達到最終的目標。表 1. 不同的模型壓縮方法。
  • 谷歌研究院出品:高效 Transformer 模型最新綜述
    眾所周知,自注意機制存在平方階的時間與存儲複雜度,這在許多情況下都阻礙了模型的可擴展性。近期,為了解決該問題,出現了各種各樣的模型變體。在下文中,我們將這類模型稱為「高效 Transformer」。  在不同的語境下,我們可以對模型的效率有不同的理 解。「效率」可能指的是模型的內存佔用,在運行模型的加速器(如 GPU)的內存時,考慮內存佔用是十分重要的。
  • 滴滴、東北大學提自動結構化剪枝壓縮算法框架,性能提升高達120倍
    滴滴 AI Labs 與美國東北大學王言治教授研究組合作,聯合提出了一種基於 AutoML 思想的自動結構化剪枝的算法框架 AutoCompress,能自動化的去尋找深度模型剪枝中的超參數,去除模型中不同層的參數冗餘,替代人工設計的過程並實現了超高的壓縮倍率。從而滿足嵌入式端上運行深度模型的實時性能需求。
  • 韓松入選,2020 IEEE智能系統十大AI青年科學家出爐
    他的研究在使用 AI 對用戶的心理模型與系統模型的協調進行可解釋交互上做出了基礎性貢獻。他的研究已經被引用了 1000 多次,並對 Explainable AI(XAI)社區產生了重大影響。他在 AAAI 2020 上編寫了關於該主題的教程。
  • 單機訓練速度提升 640 倍!獨家解讀快手商業廣告模型 GPU 訓練平臺...
    基於歷史原因,行業內推薦模型的訓練大都通過CPU來實現。然而隨著模型從Logistic Regression到深度神經網絡的演化以及硬體的發展,基於CPU的訓練系統或許已經不再是最合適的解決方案了。這些算法在很多時候可以在同樣時間內得到比使用SGD或Adagrad更好的模型。三、訓練數據分布式實時處理快手Persia的高速GPU訓練,需要大量數據實時輸入到訓練機中,由於不同模型對樣本的需求不同,對於每個新實驗需要的數據格式可能也不同。