考慮到每篇文章字數的限制,每一篇文章將按照目錄的編排包含三個小節,而且這個系列會隨著Vision Transformer的發展而長期更新。
本文目錄11 Efficient Transformer:HAT:高效的硬體感知Transformer
(來自MIT韓松團隊)
11.1 HAT原理分析
12 Efficient Transformer:Lite-Transformer:遠近注意力機制的輕量化Transformer
(來自MIT韓松團隊)
12.1 Lite-Transformer原理分析
Transformer 是 Google 的團隊在 2017 年提出的一種 NLP 經典模型,現在比較火熱的 Bert 也是基於 Transformer。Transformer 模型使用了 Self-Attention 機制,不採用 RNN 的順序結構,使得模型可以並行化訓練,而且能夠擁有全局信息。
雖然推出還不到 4 年,Transformer 已成為自然語言處理(NLP)領域裡不可或缺的一環。然而這樣流行的算法卻需要極高的算力才能實現足夠的性能,這對於受到算力和電池嚴格限制的移動端來說有些力不從心。Transformer 的高性能依賴於極高的算力,這讓移動端 NLP 嚴重受限。本文所介紹的2個工作都是 MIT 與上海交大的研究人員提出的高效的移動端 NLP 架構HAT 和 Lite Transformer,它們向在邊緣設備上部署移動級 NLP 應用邁進了一大步。
作者簡介:
HAT 和 Lite Transformer 這兩個研究都是來自MIT HAN LAB,是由 MIT 電氣工程和計算機科學系助理教授韓松領導的。韓松的研究廣泛涉足高效的深度學習模型和算法以及高效的計算機體系結構,他在斯坦福讀博期間提出的 Deep Compression 模型壓縮技術曾獲得 ICLR2016 最佳論文,論文 ESE 稀疏神經網絡推理引擎 2017 年曾獲得晶片領域頂級會議——FPGA 最佳論文獎,引領了世界深度學習加速和深度學習模型壓縮的研究,實驗室成果對業界影響深遠。
11 Efficient Transformer:HAT:高效的硬體感知Transformer (ACL 2020)論文名稱:HAT: Hardware-Aware Transformers for Efficient Natural Language Processing
論文地址:
HAT: Hardware-Aware Transformers for Efficient Natural Language Processing
https://arxiv.org/abs/2005.14187
Transformer模型在視覺和自然語言處理任務中已經無處不在,而且目前也已有了一些針對Transformer結構進行簡化的工作,但是它們都沒有考慮到不同硬體對模型結構的影響,即要想把Transformer部署在具體的硬體上面一直都是個難題。這個難題具體是:
神經網絡模型 (Transformer)計算量很密集,而硬體平臺 (微控制器,行動裝置等)往往是資源約束的 (resource-constrained),導致模型的推理過程延遲很長,極其影響用戶的體驗。同時,模型參數量過大,對部署十分不利。
比如要對一個只有30個詞的句子做機器翻譯,Transformer-Big model就需要約13 GFLOPs的計算量,在樹莓派 (Raspberry Pi)上的Latency將超過20s。
這類問題的解決方法,韓松老師實驗室--MIT HAN LAB已經有了大量探索,而且已經有了比較成熟的自動機器學習流水線(AutoML-Pipeline)來解決,具體可以參考下面的連結:
科技猛獸:解讀模型壓縮3:高效模型設計的自動機器學習流水線
https://zhuanlan.zhihu.com/p/299422728
這個自動機器學習流水線(AutoML-Pipeline)具體來說包含3步:
輸入:模型要部署的目標硬體設備。約束條件(Latency,Energy Consumption,Model Size,Memory,FLOPs等)。
輸出:部署的模型。
中間過程:
AutoML搜索最適宜的模型架構,最大化精度。
對這個模型進行壓縮,減少參數量。
對壓縮後的模型進行量化,進一步減少內存佔用。
本文介紹的HAT這種方法主要是第1步,即:藉助神經架構搜索 (NAS) 的方法,針對不同的硬體,搜索適合該硬體的最佳的Transformer結構。實驗表明,在不同硬體下,得到的Transformer結構相比其他模型更小、更快,而且不損失精度,如下圖1所示。而且,HAT這種方法可以自然地與第2步模型壓縮,第3步量化技術相結合。
圖1:藉助神經架構搜索 (NAS) 的方法,針對不同的硬體,搜索適合該硬體的最佳的Transformer結構。在把Transformer往硬體設備部署時,我們往往會忽略掉2點:
1. 計算量FLOPs不可以精確地反映模型推理的延遲Latency。
我們可以從下圖2來理解,圖2是在3種不同的硬體設備(樹莓派ARM CPU,Intel CPU和Nvidia GPU)上面得到的結果,圖2說明,在Raspberry ARM CPU上,Embedding dim的大小影響更大,但是對Intel CPU和NVIDIA GPU而言,Embedding dim就幾乎沒有影響。此外,在Raspberry ARM CPU上,FLOPS和Latency大致呈線性關係,但是在Intel CPU和NVIDIA GPU上,只有模型的層數(Layer number)呈現這種關係,對FFN的隱藏層維數hidden dim而言,3種硬體設備上相似的FLOPs會導致完全不同的Latency。這說明了,在設計Transformer的時候,如果要追求更加精細的結果,需要考慮到不同硬體的實際條件。
圖2:計算量FLOPs不可以精確地反映模型推理的延遲Latency2. 不同的硬體偏好不同的Transformer架構。
在一個硬體上優化的Transformer模型對於另一個硬體設備來說是次優的,因為Latency受不同硬體平臺上不同因素的影響。我們可以從下圖3來理解,圖3是WMT'14 En-De任務的BLEU score 以及所測量得到的延遲。可以發現針對GPU設計的 HAT (GPU) 在ARM CPU上面就不夠高效,反之亦然。
圖3:不同的硬體偏好不同的Transformer架構所以本文的初衷是藉助NAS,針對不同的硬體,搜索適合該硬體的最佳的Transformer結構。同時,將Latency作為反饋信號,反饋給搜索過程。這樣,HAT就可以搜索各種硬體的專用模型了。
整個HAT方法的流程如下圖4所示:首先我們需要構造一個很大的搜索空間,叫做SuperTransformer,它包括了design space中的所有candidates。作者將它訓練到收斂,通過權值共享,就能很方便地得到大量的SubTransformer。最後執行Evolutionary Search,並把神經網絡模型在目標硬體上的Latency直接作為反饋信號搜索出適合這個目標硬體設備的最佳Transformer。從SuperTransformer中採樣子網絡,不同子網絡中相同的部分共享權重 (Weight Sharing)。如此一來,就不必為每個自網絡訓練單獨的模型,這就大大加快了搜索速度。最後從頭訓練搜索得到的SubTransformer。通過實驗顯示,本文搜索出的模型結構不僅更小更快,而且在效果上甚至更好,超過了Evolved Transformer等結構。
圖4:HAT方法的流程HAT具體做法-搜索空間:NAS方法最重要的問題是搜索空間的設計,HAT在設計搜索空間時打破了Transformer 結構設計的2條常規,分別如圖5和7所示,即:
所有 decoder layers 僅僅關注最後一個 encoder layer。Arbitrary Encoder-Decoder Attention
先看第1條:解碼器的每一層同時關注編碼器的多層而不僅僅是最後一層,作者將其命名為Arbitrary Encoder-Decoder Attention,如圖5所示。常規Transformer的設計中,所有 decoder layers 僅僅關注最後一個 encoder layer的輸出,這就形成了一個信息瓶頸 (information bottleneck),如下圖5左側所示,強制所有的 decoder layers 只學習那些高級的特徵,而忽略大量的低級別的信息。Arbitrary Encoder-Decoder Attention的目的就是學習編碼器和解碼器之間最合適的連接,如下圖5右側所示。該機制沒有引入任何額外的參數,使得每個 decoder layer 可以選擇多個 encoder layer。
圖5:HAT打破常規1:All decoder layers only attend to the last encoder layer下圖6進一步闡釋了Arbitrary Encoder-Decoder Attention的機制。一個 decoder layer 中的每個attention可以關注來自多個 encoder layer 的輸出,充分利用 encoder layer 提取的特徵。來自 encoder layer 的Key和Value在sentence length這個維度被concat在一起,送到cross-attention模塊中。
圖6:Arbitrary Encoder-Decoder Attention機制。Heterogeneous Transformer Layers
再看第2條:所有layers可以不是identical的,如下圖7所示。以前的Transformer對所有layers重複一個架構。但對HAT的Transformer而言,可以變化的比較重要的參數有:Embedding dim,Hidden dim,Head數,Layer數,再加上第1條解碼器的每一層同時關注編碼器的多層而不僅僅是最後一層。以上就定義了整個搜索空間。
這裡解釋下Embedding dim 和 Hidden dim 的區別:
我們知道Q, K, V可以看做是矩陣,其第1維是Sequence length,代表序列的長度。第2維就是Embedding dim ,代表序列中單個向量的長度。
Hidden dim來自FFN層,我們知道FFN會先把[N, Embedding dim]的矩陣變成[N, expansion ratio × Embedding dim]的矩陣,那這個Hidden dim = expansion ratio × Embedding dim。
圖7:HAT打破常規2:All the layers are identical這些超參數的具體設置為:
Embedding dim設置為:[512, 640]
hidden dim設置為:WMT task:[1024, 2048, 372]; IWSLT task:[2048, 1024, 512]
head number設置為:WMT task:[4, 8]; IWSLT task:[4, 2]
decoder layers設置為:[1, 2, 3, 4, 5, 6]
encoder layers統一設置為:[6],因為encoder的Latency只佔總Latency的不到5%。
這樣定義搜索空間之後,總的SubTransformer數量約為個。
所以,如果要訓練所有的模型並且來評估它們的BLEU值,顯然是不可能的。於是,本文提出SuperTransformer——作為所有可能的模型結構的一個「父類」——來囊括所有可能的模型。SuperTransformer是搜索空間中「最大」的一個模型,所有的其他模型都是SuperTransformer中的一個子網絡,並且共享相同的部分。一個子網絡只需要取SuperTransformer中的前端一部分參數即可。
下圖是一個示例。藍色部分表明:如果搜索空間中Embedding維數最大為640,當前想要測評的模型維數是512,那麼我們只需要取整個大矩陣的前512維。橙色部分同理:如果搜索空間中Hidden維數最大允許是1024,而當前考慮的模型是768,也只需要取前768維即可。
圖8:SuperTransformer的權重共享,一個子網絡只需要取SuperTransformer中的前端一部分參數即可。在SuperTransformer訓練中,對所有可能的SubTransformer進行統一採樣,並更新相應的權重。在實踐中,SuperTransformer只需要像普通Transformer模型一樣接受相同步驟的訓練,速度快,成本低。經過訓練後,我們可以直接在驗證集上評估所有的SubTransformer的性能。
HAT具體做法-搜索策略:給定一個目標硬體與它對應的Latency要求,我們執行Evolutionary Search來找到一個令人滿意的SubTransformer。評估SubTransformer的硬體延遲有兩種方法:
在線方法,即在搜索過程中測量模型的Latency。離線方法,我們訓練一個延遲預測器 (Latency Predictor) 來提供延遲。作者在實做中使用了離線方法。因為在搜索過程中,我們需要評估所有SubTransformer的性能,那麼使用真實數據對每個模型都測量它們的Latency這一過程將大量耗時。基於此,作者提出使用Latency Predictor,這也是MIT HAN LAB的常用套路,在之前的HAQ,APQ中也有使用。即直接訓練一個小的神經網絡,輸入模型結構,預測該模型在此硬體上的Latency。
問:延遲預測器 (Latency Predictor) 的輸入和輸出是什麼?
答:輸入是:
1. Encoder layer number, 2. Encoder Embedding dim,3. Encoder hidden dim
dim,4. Encoder average self-attention heads,
5. Decoder layer number,6. Decoder Embedding dim,7. Decoder hidden
dim,8. Decoder average self-attention heads,
9. average encoder-decoder attention heads,10. average number of encoder
layers each decoder layer attends (每個decoder層關注的encoder層數量的均值)。
輸出是:Predicted Latency。
問:延遲預測器 (Latency Predictor) 的結構是什麼?
答:MIT HAN LAB的很多文章都涉及了Latency Predictor,結構一般是3層的MLP,它的hidden dim本文為400,激活函數為ReLU。為什麼用3層而不是1層,5層?因為1層預測不準,3層效果已經很好,超過3層將對準確率提升意義不大,白白增加模型參數。
乍一看這個方法似乎不靠譜,但是實際的效果卻很好。如下圖9所示。網絡預測得到的Latency和實際的幾乎沒有差別,而我們需要做的,只是從每個硬體中採集真實的訓練數據。本文為每個硬體提供了2000個樣本,按照8:1:1的方法劃分數據集。Latency Predictor只用在搜索過程中,在最後實驗的時候還是用的真實測得的Latency。
圖9:Latency Predictor很精確作者使用進化算法來進行搜索過程。如下圖10所示。Search Engine僅將Latency小於目標硬體Latency約束的SubTransformer添加到進化算法的種群Population中。進化算法的超參數設置為:一共進行30輪,種群數(population)設置為125,parents為25,mutation population為50,crossover population也為50。
圖10:使用進化算法,聯合SuperTransformer和Latency Predictor得到適合目標硬體設備的最優模型下圖是HAT的方法概覽:
第1步就是定義搜索空間,之後訓練一個SuperTransformer並從中採樣子網絡。
第2步是搜集實際目標硬體Hardware-Latency的數據。
第3步把這個數據訓練Latency Predictor。
第4步是藉助Latency Predictor,使用進化算法搜索得到最佳的模型結構。
第5步是重新訓練搜索得到的模型,得到最終的Performance。
圖11:HAT方法概覽總的來說本文的貢獻有3點:
基於不同硬體搜索不同的Transformer結構,和硬體實際條件更加契合。針對不同的硬體平臺,可以得到低延遲推理的專用模型。(基於不同硬體搜索不同的網絡架構其實在MIT HAN LAB之前已有在分類任務上實現。)提出SuperTransformer搜索範式,大大加快搜索速度 (創新點在arbitrary encoder-decoder attention和heterogeneous layer,搜索策略不新奇,依舊是One-Shot NAS)。得到了一些設計Transformer的啟發,即:GPU更喜歡shallow 和 wide的模型,而ARM CPU更喜歡deep 和 thin的模型。Experiments:
研究者在四個機器翻譯任務和三種硬體平臺上進行了實驗和分析。四個任務為: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, Levenshtein Transformer, Evolved Transformer,和 Lite Transformer。
下圖是在En-De和En-Fr上,不同模型的Latency和BLEU值結果。比如最左側的2個圖是在Raspberry Pi 的 ARM CPU上,HAT比Transformer Big快3.0倍,並且小3.7倍多,甚至效果還更好;在Intel CPU、Nvidia GPU上也有不小的提升。在Nvidia GPU上,藍色的點對應的是dimension scaling的結果,我們可以看到其連成的線段已經近乎垂直,這說明想靠著dimension scaling的方式減少模型的延遲是根本行不通的。
圖12:En-De和En-Fr上,不同模型的Latency和BLEU值結果下圖13為HAT模型和原版Transformer,Evolved Transformer參數量,延遲,計算量,BLEU得分,訓練的GPU時長,排放溫室氣體量以及雲計算成本等各項指標的對比。相比於Baseline模型,HAT取得了大幅的加速以及模型參數的縮減。同時在訓練時長方面,因為採用了One-Shot NAS的辦法,所以優勢巨大,因此導致了溫室氣體的排放量以及雲計算成本的下降。
圖13:參數量,延遲,計算量等不同指標的對比專用模型的必要性
兩個搜索得到的模型表現如下圖14所示。可以看到,GPU 高效模型在 ARM CPU 上並不是最高效的,ARM CPU 高效模型在 GPU 上也不高效 (The efficient model for GPU is not efficient for ARM CPU and vice versa)。這進一步證明了為不同硬體設計專用模型的必要性。
圖14:兩個搜索得到的模型表現與 Evolved Transformer 對比
在樹莓派上運行 WMT' 14 EN-Fr (英語 - 法語) 任務時,相比 Evolved Transformer,HAT 可以取得 2.7 倍加速,3.7 倍模型壓縮,3.2 倍計算量降低,並節省超過 1 萬倍的搜索開銷。
圖15:與 Evolved Transformer 各種指標的對比與 Levenshtein Transformer 和 Lite Transformer 對比
下圖16為不同模型在樹莓派上運行 WMT' 14 EN-De (英語 - 德語) 任務的對比。相比其他模型,HAT 可以獲得最高的 BLEU 和最低的運行時間。值得注意的是,HAT 方法與 Levenshtein 和 Lite Transformer 中提出的新操作具有正交性 (Orthogonal),可以結合使用。
圖16:不同模型在樹莓派上運行 WMT' 14 EN-De (英語 - 德語) 任務的對比與其他壓縮方法相兼容
下圖17為在 WMT』14 EN-Fr (英語 - 法語) 任務上,對 HAT 搜索模型進行量化的結果,在 4-bit 量化的情況下,HAT 模型與基線 Transformer 有相近的精度,但是模型大小可壓縮 25 倍。
圖17:HAT與其他壓縮方法相兼容不同硬體上得到的最好的結構
最後來看看使用HAT在不同硬體上得到的最好的結構是怎樣的。如下圖18所示,可以發現,對Raspberry Pi ARM CPU和NVIDIA GPU,HAT得到的模型結構是很不同的。二者在層數、維度數和Attention方式上都有差別。Raspberry Pi上搜索得到的模型是deep and thin,但GPU上搜索得到的模型是shallow and wide。GPU Latency 對 Embedding dim 和 hidden dim 不敏感,但Raspberry Pi對其敏感度很高。所以這可以用來指導手工設計:在GPU上,我們可以減少layers的同時增加Embedding dim 和 hidden dim,這樣一方面能夠降低Latency,另一方面也能保持高性能。
圖18:HAT在不同硬體上得到的最好的結構12 Efficient Transformer:Lite-Transformer:遠近注意力機制的輕量化Transformer (ICLR 2020)論文名稱:Lite Transformer with Long-Short Range Attention
論文地址:
Lite Transformer with Long-Short Range Attention
https://arxiv.org/abs/2004.11886
12.1 Lite-Transformer原理分析:Transformer模型因其訓練效率高、捕獲長距離依賴能力強等特點,已經在自然語言處理中得到廣泛應用。在此基礎上,現代最先進的模型,如BERT,能夠從未標註的文本中學習強大的language representation,甚至在一些很有挑戰性的問答任務上超越人類。但它需要大量計算去實現高性能,比如一個Transformer模型翻譯一個長度不超過30個單詞的句子需要大約10G 的Mult-Adds。而這不適合受限於硬體資源和電池嚴格限制的移動應用,比如智慧型手機,手環,物聯網設備等。那麼如何減少Transformer的計算量呢?看了上面的HAT我們知道一種辦法是通過減少Embedding size 。但是這存在的一個問題是:這樣做在減少計算量的同時也削弱了Transformer捕獲長距離和短距離關係的能力。
Lite-Transformer這項研究提出了一種高效的模塊 —— LSRA,其核心是長短距離注意力(Long-Short Range Attention,LSRA),其中一組注意力頭(通過卷積)負責局部上下文建模,而另一組則(依靠注意力)執行長距離關係建模。
這樣的專門化配置使得模型在機器翻譯、文本摘要和語言建模這3個語言任務上都比原版 transformer 有所提升,基於LSRA所構建的Lite Transformer達到了行動裝置計算量所要求的500M Mult-Adds。以WMT 2014 English-German任務為例,在計算量限制為500M Mult-Adds或者100M Mult-Adds時,Lite Transformer的性能比原版 Transformer 的 BLEU 值比分別比 transformer 高 1.2或1.7。結合剪枝和量化技術,研究者進一步將 Lite Transformer 模型的大小壓縮到原來的 5%。
對於語言建模任務,在大約 500M MACs 上,Lite Transformer 比 transformer 的困惑度低 1.8。值得注意的是,對於移動 NLP 設置,Lite Transformer 的 BLEU 值比基於 AutoML 的Evolved Transformer 高 0.5,而且AutoML方法所需要的搜索算力超過了250 GPU years,這相當於5輛汽車的終身碳排放量。
Lite Transformer具體方法:
我們將用於文本處理的Self-attention稱為1-D attention,用於圖像識別的Self-attention稱為2-D attention,用於視頻處理的Self-attention稱為3-D attention。首先看看Self-attention的計算複雜度,如下圖19所示:
圖19:Self-attention的計算複雜度圖19中的 代表Sequence的長度, 代表Embedding dim。所以Self-attention的計算複雜度為: 。這步近似是因為通常對於圖像或者視頻來講, 通常很大。比如說16幀112 × 112的視頻, 。
這對於整個Transformer模型的計算量來說有什麼壞處?
因為Transformer模型裡面的卷積層和FFN的計算量與 呈線性關係,而Transformer模型裡面的Self-attention的計算量與 呈平方關係(如上面推導)。所以整個Transformer模型的計算量當 增大時也同樣會變得很大。
如何解決這個問題呢?
首先 一般是不能改變的,因為句子,圖片或者視頻的大小一般是根據數據集定的,不能變化。那麼現在能變化的是Embedding dim,也就是這個 。那降低這個Embedding dim會帶來什麼問題呢?
我們先從另一個角度理解Transformer的結構,如下圖20所示。圖20上面這個圖代表Transformer的1個block,現在我們換個角度,用下面這個圖的劃分方式來理解。每個block先通過FFN降低維度,再經過Self-attention處理,再通過FFN升高維度。我們把這種設計成為bottleneck design。所以我們可以發現傳統Transformer結構的一個非常不好的地方,即:如圖21所示,傳統的Transformer結構把大量的計算都花費在了FFN結構身上,而且,FFN結構的計算複雜度是 而,Self-attention的計算複雜度是 。當 越小時,FFN結構所佔用的計算複雜度的百分比就越大,對於文本翻譯這類任務, 一般取幾十;對於圖像處理這類任務, 一般取幾百。所以這種情況對於文本翻譯這類任務尤其嚴重。
那麼把大量的計算都花費在了FFN結構身上有什麼壞處呢?
不論是文本處理還是圖像處理的任務,都需要對序列上下文的關係進行建模,而這個建模的工作是靠Self-attention來完成的。但如果我們像上面預設的減少Embedding size ,就會使得Self-attention layer的feature dimension減得很小,勢必會影響Self-attention layer的性能。而FFN layer卻佔用了大量的計算量,這就是傳統的Transformer結構的不合理之處,即:由於在NLP任務中通常 較小,導致有用的Self-attention layer所佔用的計算量比例很小,減少Embedding size會嚴重影響Self-attention layer的性能,使得我們無法在保證性能的前提下大幅減少計算量。
圖20:從另1個角度理解Transformer的結構圖21:傳統的Transformer結構把大量的計算都花費在了FFN結構身上所以作者認為,bottleneck design 對於1-D Self-attention layer來說不是最優的,所以作者首先設計了如下圖22所示的Flattened Transformer Block,它使得特徵在進入Self-attention layer之前不進行降維,使得attention layer佔據了絕大部分計算量。這樣的好處是為進一步優化留出更大的空間。作者還在WMT' 14 En-Fr數據集上測試了性能,可以在稍大一點的計算中獲得相當的性能,而計算量可以通過LSRA進行大幅降低而不影響性能。
對比一下之前的方法:
之前想通過減少Embedding dim來降低計算量,但是由於 bottleneck design 的缺點,使得Self-attention受到了嚴重的影響,影響了模型的性能。
現在通過減少LSRA來降低計算量,由於 Flattened Transformer Block,使得計算量可以通過LSRA進行大幅降低而不影響性能。
圖22:Flattened Transformer Block使得attention layer佔據大部分計算量如下圖23所示為Self-attention layer的可視化,每個格子代表原句的每個單詞和目標句子的每個單詞之間的關係權重大小。顏色越深表示關係越強。隨著權重的增加(顏色加深),源句中的第 個詞更加注意目標句中的第 個詞。這個可視化圖有2個特點,即:稀疏和對角線。稀疏,如圖中的藍色虛線部分所示;對角線,如圖中的黃色虛線部分所示。即:Self-attention layer同時捕捉到了稀疏和對角線這2種模式化特徵。它們代表了一些特定單詞之間的關係:稀疏表示長距離信息間的關係,對角線表示近距離信息間的關係。研究者將前者稱為全局 (global) 關係,將後者稱為局部 (local) 關係。
圖23:Self-attention layer的可視化那麼Self-attention同時捕獲長距離信息間的關係和近距離信息間的關係,這看似很合理,但是實際上卻並不高效。為什麼呢?因為這使得Self-attention足夠靈活 (Flexible),但不夠專門化 (Specialized)。
這句話是什麼意思?我們以硬體設計為例,CPU是一種非常general-purpose的硬體設備,換句話講,CPU非常靈活;而FPGA是一種非常specialized的設備。但是,CPU 等通用硬體的效率卻比 FPGA 等專用硬體低。所以作者認為,看似非常靈活的Self-attention實際上是非常低效的,而我們需要的是足夠專門化的模塊,以得到高效率。
通俗地講,我們希望Self-attention這個模塊更加專門化一點。為了解決該問題,該研究提出一個更專門化的架構,即長短距離注意力 (LSRA)。
圖24:以硬體設計為例解釋Flexible與Specialized的優缺點長短距離注意力 (LSRA)哪裡專門化呢?在翻譯任務中,注意力模塊必須捕獲全局和局部上下文信息。LSRA 模塊遵循兩分支設計,如下圖25所示。左側注意力分支負責捕獲全局上下文,右側卷積分支則建模局部上下文。研究者沒有將整個輸入饋送到兩個分支,而是將其沿通道維度分為兩部分,然後由後面的 FFN 層進行混合。這種做法將整體計算量減少了 50%。
圖25:LSRA模塊左側分支處理全局信息:是正常的Self-attention模塊,不過通道維度減少了一半。
右側分支處理局部關係:一個自然的想法是對序列應用卷積。為了進一步減少計算量,研究者將普通卷積替換為輕量級的版本,該版本由線性層linear layers和Depth-wise convolution組成。
這裡的卷積為1維的Depth-wise Convolution 。
這裡的GLU模塊出自論文:Language Modeling with Gated Convolutional Networks,下面簡單介紹下GLU模塊:
目前語言模型主要基於RNN,這篇文章提出了一種新穎的語言模型,仿照LSTM中的門限機制,利用多層的CNN結構,每層CNN都加上一個輸出門限。
圖26:GLU模型GLU模型:
公式中的 是上一層的輸出向量 (或者初始輸入的詞序列向量), ,,。 為詞序列的長度, 的詞向量的維度, 為卷積核的大小, , 為偏置。最後每層的輸出每個詞的向量表示:。( 為輸入, 為模型的層數)。最後把 個層stack起來以後,再乘以linear transformation 和softmax 即得到上下文每個詞的representation 。
總之,通過這種方式,研究者將注意力模塊和卷積模塊並排放置,引導它們對句子進行全局和局部的不同角度處理,從而使LSRA模塊更加專門化,以求實現更高的效率。
下圖27為原版Transformer模型以及使用LSRA模塊得到的Transformer模型的attention weights的可視化結果。顯然,LSRA中的Self-attention模塊不試圖對全局上下文和局部上下文都進行建模,而是只關注全局上下文捕獲(因為沒有對角線),而將局部上下文捕獲留給卷積分支。
圖27:原版Transformer模型以及使用LSRA模塊得到的Transformer模型的attention weights的可視化結果總的來說本文的貢獻有3點:
發現bottleneck design的結構對於1-D attention (文本處理) 來說不是最優的。提出一種多分支的特徵提取器 Long-Short Range Attention (LSRA),其中卷積操作幫助捕捉局部上下文,而attention用來捕捉全局上下文。基於LSRA所構建的Lite Transformer達到了行動裝置計算量所要求的500M Mult-Adds,在3種任務上獲得了一致的性能提升,與AutoML-based方法Evolved Transformer相比也獲得了性能的提升,並大大減少訓練成本。Experiments:
首先我們介紹下Transformer模型的設計目標,具體來說就是計算量和模型大小的目標值應該怎麼選擇。
計算量目標:一個ARM Cortex-A72 mobile CPU的算力大概是48 GFLOPs,如果想要每秒完成50個句子的翻譯,那麼Transformer模型翻譯每個句子的計算量上限就是960M FLOPs (480M Mult-Adds)。所以我們把機器翻譯任務的計算量要求設置為:翻譯一個長度為30個單詞的句子的計算量低於500M Mult-Adds (or 1G FLOPs)。
參數量目標:MobileNet 設置為7M parameters限制,這裡設置為近似值10M。
數據集和評估
作者在機器翻譯、文本摘要和語言建模三個任務上進行了實驗和評估。
具體而言,機器翻譯任務使用了三個基準數據集:IWSLT' 14 德語 - 英語 (De-En)、WMT 英語 - 德語 (En-De)、WMT 英語 - 法語(En-Fr)。
文本摘要任務使用的是 CNN-DailyMail 數據集。
語言建模任務則在 WIKITEXT-103 數據集上進行。
機器翻譯實驗結果:
IWSLT 實驗結果
下圖28為Lite Transformer 在 IWSLT' 14 De-En 數據集上的定量結果。並與 transformer 基線方法和 LightConv 做了對比。在大約 100M Mult-Adds 時,Lite Transformer 模型的 BLEU 值比 transformer 高出 1.6。
圖28:IWSLT'14 De-En 數據集結果WMT 實驗結果
下圖29為Lite Transformer 在 WMT' 14 En-De and WMT' 14 En-Fr 數據集上的定量結果。並與 transformer 基線方法做了對比。Lite Transformer在總計算量和模型參數量之間實現了更好的平衡。在大約 100M Mult-Adds 時,Lite Transformer 模型的 BLEU 值比 transformer 分別高出了 1.2和1.7;在大約 300M Mult-Adds 時,Lite Transformer 模型的 BLEU 值比 transformer 分別高出了 0.5和1.5。
圖29:WMT' 14 En-De and WMT' 14 En-Fr 數據集結果WMT En-Fr數據集實驗結果的trade-off曲線如下圖30所示。
圖30:WMT En-Fr數據集實驗結果的trade-off曲線與其他壓縮方法相兼容
與HAT方法一樣,Lite Transformer方法也與其他幾個經典的模型壓縮方法正交,例如剪枝和量化,如圖31所示。作者使用K-means將模型權重量化為8位,並根據每一層的敏感度對模型進行剪枝。利用這兩種模型壓縮技術,作者實現了18.2倍模型大小的壓縮,而BLEU分數的下降可以忽略不計。
圖31:Lite Transformer方法也與其他幾個經典的模型壓縮方法正交與 Evolved Transformer 對比
相比 Evolved Transformer,在大約 100M Mult-Adds 時,Lite Transformer 模型的 BLEU 值比 Evolved transformer 高出了 0.5;在大約 300M Mult-Adds 時,Lite Transformer 模型的 BLEU 值比 Evolved transformer 高出了 0.2。
圖32:與 Evolved Transformer 各種指標的對比總結:本文介紹了 MIT HAN LAB 提出的2種構建輕量級Transformer的方法。HAT 通過NAS方法,結合延遲預測器,搜索Super Transformer的最優解,Lite-Transformer通過設計輕量化的長短距離注意力 (LSRA) 模塊,遵循兩分支設計,左側注意力分支負責捕獲全局上下文,右側卷積分支則建模局部上下文,減少整體的計算量。在設計輕量化的Transformer的時候,一般的想法是從 Embedding dim,Hidden dim,Head數,Layer數,解碼器關注的編碼器位置 這五個維度減少Transformer的model size,必要時可以從諸如LSRA的輕量化結構入手。