模型壓縮95%:Lite Transformer,MIT韓松等人

模型壓縮95%:Lite
Transformer,MIT韓松等人

Lite Transformer with Long-Short Range Attention

Zhanghao Wu, Zhijian Liu, Ji Lin, Yujun Lin, Song Han

· 論文地址:https://arxiv.org/abs/2004.11886v1

· GitHub 地址:https://github.com/mit-han-lab/lite-transformer

摘要

Transformer在自然語言處理(如機器翻譯、問答等)中已經非常普遍,但是它需要大量的計算才能達到較高的性能,這使得它不適合於硬件資源和電池的嚴格限制下的移動應用。本文提出了一種高效的移動NLP架構Lite Transformer,以便於在邊緣設備上部署移動NLP應用。關鍵原語是長-短距離注意(LSRA),其中一組負責人專門研究局部上下文建模(通過卷積),另一組負責人專門研究長距離關係建模(通過注意)。這種專業化在三個公認的語言任務(機器翻譯、抽象摘要和語言建模)上比普通的轉換程序帶來了一致的改進。在資源受限(500米/100米MACs)的情況下,Lite Transformer的性能分別比WMT的14英法文高1.2/1.7 BLEU。Lite-Transformer在BLEU評分下降0.3的情況下,將變壓器基礎模型的計算量減少了2.5x。結合剪枝和量化,進一步壓縮了Lite-Transformer的模型尺寸18.2x,在語言建模方面,Lite-Transformer在500M左右的MACs下比Transformer的複雜度降低了1.8。值得注意的是,Lite Transformer比基於AutoML的演進Transformer在移動NLP設置上的BLEU高0.5,而無需花費超過250 GPU年的昂貴架構搜索。

Transformer 的高性能依賴於極高的算力,這讓移動端 NLP 嚴重受限。在不久之前的 ICLR 2020 論文中,MIT 與上海交大的研究人員提出了一種高效的移動端 NLP 架構 Lite Transformer,向在邊緣設備上部署移動級 NLP 應用邁進了一大步。

雖然推出還不到 3 年,Transformer 已成爲自然語言處理(NLP)領域裏不可或缺的一環。然而這樣流行的算法卻需要極高的算力才能實現足夠的性能,這對於受到算力和電池嚴格限制的移動端來說有些力不從心。

在 MIT 最近的研究《Lite Transformer with Long-Short Range Attention》中,MIT與上海交大的研究人員提出了一種高效的移動端 NLP 架構 Lite Transformer,向在邊緣設備上部署移動級 NLP 應用邁進了一大步。該論文已被人工智能頂會 ICLR 2020 收錄。

該研究是由 MIT 電氣工程和計算機科學系助理教授韓松領導的。韓松的研究廣泛涉足深度學習和計算機體系結構,他提出的 Deep Compression 模型壓縮技術曾獲得 ICLR2016 最佳論文,論文 ESE 稀疏神經網絡推理引擎 2017 年曾獲得芯片領域頂級會議——FPGA 最佳論文獎,引領了世界深度學習加速研究,對業界影響深遠。
在這裏插入圖片描述
圖1:左:最近NLP模型的大小增長迅速,超過了移動限制很多。右圖:基於AutoML的NLP模型的搜索成本高得令人望而卻步,因爲它會排放碳,二氧化碳幾乎是汽車壽命平均排放量的5倍。
在這裏插入圖片描述
在這裏插入圖片描述
注意機制在各種應用中得到了廣泛的應用,包括1-Dlanguage處理(Vaswani等人,2017年)、二維(圖像識別)和三維(視頻識別)(Wang等人2018年)。它計算所有輸入元素之間的成對點積,以對兩個短期元素進行建模以及長期的關係。儘管該操作有效,但它引入了大量計算。

假設元素的數量(例如,語言處理中標記的長度、像素的數量在圖像等中,輸入到注意層的是N,而特徵(即通道)的維數是d,則點積所需的計算是N2d。

爲了解決這個難題,通常的做法是首先使用注意力之前的線性投影層,然後增加尺寸(如圖2所示)。在transformer的原始設計中(Vaswani等人,2017),通道尺寸注意模塊比FFN層小4倍。同樣,在非本地視頻中,網絡(Wang等人,2018),在應用非本地注意模塊。這種做法節省了16×或4×的計算量。然而,它也減少了語境對特徵維度較小的注意層的捕獲能力。情況可能更糟的是語言處理,因爲注意力是獲取上下文的主要模塊(不像卷積進行圖像和視頻的主要信息捕獲)。

Transformer 在自然語言處理任務(如機器翻譯、問答)中應用廣泛,但它需要大量計算去實現高性能,而這不適合受限於硬件資源和電池嚴格限制的移動應用。

這項研究提出了一種高效的移動端 NLP 架構——Lite
Transformer,它有助於在邊緣設備上部署移動 NLP 應用。其核心是長短距離注意力(Long-Short Range Attention,LSRA),其中一組注意力頭(通過卷積)負責局部上下文建模,而另一組則(依靠注意力)執行長距離關係建模。

這樣的專門化配置使得模型在三個語言任務上都比原版 transformer 有所提升,這三個任務分別是機器翻譯、文本摘要和語言建模。

在資源有限的情況下(500M/100M MACs),Lite Transformer 在 WMT’14 英法數據集上的BLEU 值比分別比 transformer 高1.2/1.7。Lite Transformer 比transformer base 模型的計算量減少了 60%,而 BLEU 分數卻只降低了 0.3。結合剪枝和量化技術,研究者進一步將 Lite Transformer 模型的大小壓縮到原來的 5%。

對於語言建模任務,在大約 500M MACs 上,Lite Transformer 比 transformer 的困惑度低1.8。
在這裏插入圖片描述
值得注意的是,對於移動 NLP 設置,Lite Transformer 的 BLEU 值比基於 AutoML 的 Evolved Transformer 高 0.5,而且它不需要使用成本高昂的架構搜索。

從 Lite Transformer 與 Evolved Transformer、原版 transformer 的比較結果中可以看出,Lite Transformer 的性能更佳,搜索成本相比 Evolved Transformer 大大減少。
在這裏插入圖片描述
那麼,Lite Transformer 爲何能夠實現高性能和低成本呢?接下來我們來了解其核心思想。

長短距離注意力(LSRA)

NLP 領域的研究人員試圖理解被注意力捕捉到的上下文。Kovaleva 等人 (2019) 和 Clark 等人(2020) 對 BERT 不同層的注意力權重進行了可視化。

如下圖 3b 所示,權重 w 表示源句單詞與目標句單詞之間的關係(自注意力也是如此)。隨着權重 w_ij 的增加(顏色加深),源句中的第 i 個詞更加註意目標句中的第 j 個詞。注意力圖通常有很強的模式化特徵:稀疏和對角線。它們代表了一些特定單詞之間的關係:稀疏表示長距離信息間的關係,對角線表示近距離信息間的關係。研究者將前者稱爲「全局」關係,將後者稱爲「局部」關係。
在這裏插入圖片描述
圖 3:Lite Transformer 架構 (a) 和注意力權重的可視化。傳統的注意力 (b) 過於強調局部關係建模(參見對角線結構)。該研究使用卷積層專門處理局部特徵提取工作,以高效建模局部信息,從而使注意力分支可以專門進行全局特徵提取 ©。

在翻譯任務中,注意力模塊必須捕獲全局和局部上下文,這需要很大的容量。與專門化的設計相比,這並非最佳選擇。以硬件設計爲例,CPU 等通用硬件的效率比 FPGA 等專用硬件低。研究者認爲應該分別捕捉全局和局部上下文。模型容量較大時,可以容忍冗餘,甚至可以提供更好的性能。但是在移動應用上,由於計算和功率的限制,模型應該更加高效。因此,更需要專門化的上下文捕獲。

爲了解決該問題,該研究提出一個更專門化的架構,即長短距離注意力(LSRA),而不是使用處理 “一般” 信息的模塊。該架構分別捕獲局部和全局上下文。

如圖 3a 所示,LSRA 模塊遵循兩分支設計。左側注意力分支負責捕獲全局上下文,右側卷積分支則建模局部上下文。研究者沒有將整個輸入饋送到兩個分支,而是將其沿通道維度分爲兩部分,然後由後面的 FFN 層進行混合。這種做法將整體計算量減少了 50%。

左側分支是正常的注意力模塊(Vaswani et al. (2017)),不過通道維度減少了一半。至於處理局部關係的右分支,一個自然的想法是對序列應用卷積。使用滑動窗口,模塊可以輕鬆地覆蓋對角線組。爲了進一步減少計算量,研究者將普通卷積替換爲輕量級的版本,該版本由線性層和深度卷積組成。通過這種方式,研究者將注意力模塊和卷積模塊並排放置,引導它們對句子進行全局和局部的不同角度處理,從而使架構從這種專門化設置中受益,並實現更高的效率。

實驗設置

數據集和評估 研究者在機器翻譯、文本摘要和語言建模三個任務上進行了實驗和評估。

具體而言,機器翻譯任務使用了三個基準數據集:IWSLT’14 德語 - 英語 (De-En)、WMT 英語 - 德語 (En-De)、WMT 英語 - 法語(En-Fr)。

文本摘要任務使用的是 CNN-DailyMail 數據集。

語言建模任務則在 WIKITEXT-103 數據集上進行。

架構

模型架構是基於序列到序列學習的編碼器 - 解碼器。在機器翻譯任務中,針對 WMT 數據集,基線模型基於 Vaswani 等人提出的模型。對於 IWSLT 數據集,基線模型遵循 Wu 等人的設置。對於文本摘要任務,研究者採用了與 WMT 相同的模型。至於語言建模任務,模型與 Baevski & Auli(2019) 一致,但模型尺寸較小。

該研究提出的架構首先將 transformer base 模型中的 bottleneck 拉平,然後用 LSRA 替換自注意力。更具體地說,是使用兩個專門的模塊,一個注意力分支和一個卷積分支。

實驗結果

機器翻譯

表 1 展示了 Lite Transformer 在 IWSLT’14 De-En 數據集上的定量結果,並與 transformer 基線方法和 LightConv 做了對比。在大約 100M Mult-Adds 時,Lite Transformer 模型的 BLEU 值比 transformer 高出 1.6。
在這裏插入圖片描述
表 1:IWSLT’14 De-En 上的結果。

表 2 中的定量結果表明,在 100M Mult-Adds 設置下,Lite Transformer 在 WMT En-De 數據集和 WMT En-Fr 數據集上的 BLEU 值分別比 Transformer 高出 1.2 和 1.7,在 300M Mult-Adds 設置下,也有 0.5 和 1.5 分的提升。
在這裏插入圖片描述
表 2:在 WMT’14 En-De 和WMT’14 En-Fr 上的結果。

研究者還提供了模型在 WMT En-Fr 上的權衡曲線,如圖 4a 所示,Lite Transformer 一直優於原版 transformer。
在這裏插入圖片描述圖 4:在 WMT En-Fr 數據集上的機器翻譯權衡曲線,以及在 WIKITEXT-103 數據集上的語言建模權衡曲線。兩個曲線都說明了在移動設置下,Lite Transformer 比 transformer 性能更佳(藍色區域)。

與自動化設計模型的對比

與基於 AutoML 的 Evolved Transformer(ET)相比,Lite Transformer 在移動設置中也有明顯的改進。此外,在 100M 和 300M 的Mult-Adds 下,Lite Transformer 的BLEU 值分別比 ET 高 0.5 和 0.2,詳見表 3。
在這裏插入圖片描述
表 3:不同 NMT 模型的性能和訓練成本。

文本摘要
在這裏插入圖片描述
表 4:在 CNN-DailyMail 數據集上的文本摘要結果。
在這裏插入圖片描述
表 5:在 WIKITEXT-103 數據集上的語言建模結果。

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章